You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@skywalking.apache.org by wu...@apache.org on 2026/10/08 12:13:54 UTC
(skywalking) branch master updated: Fix what the confirming reviews of AI agent withholding found (#14124)
This is an automated email from the ASF dual-hosted git repository.
wu-sheng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/skywalking.git
The following commit(s) were added to refs/heads/master by this push:
new 49c575adec Fix what the confirming reviews of AI agent withholding found (#14124)
49c575adec is described below
commit 49c575adecc7e5c6993c32409355851f495647c2
Author: 吴晟 Wu Sheng <wu...@foxmail.com>
AuthorDate: Thu Oct 8 20:13:48 2026 +0800
Fix what the confirming reviews of AI agent withholding found (#14124)
The reviews that confirmed the fixes of #14123 found more places where
the OAP's withholding did not hold, or did not agree with the
Sessionizer's own rules.
- The masking walk stopped at a set of tools, so a prompt key inside one
kept its text. A set of tools is now masked as every value is: what it
is as a whole first, a system message or a list of messages with a
system pair, then each tool, a list inside the set as more of its
tools, and anything else in it like any other member. A masked tool is
not walked again, so a tool named like a prompt key keeps the tools'
marker.
- A body that is not UTF-8 text was written as a literal under its
stored digest and size, so the served body no longer matched its own
manifest. Such a body cannot be a literal byte for byte, and is now
withheld whole; so is a request that is not one JSON object in UTF-8
text.
- A body is now held as the Session Data page says a reader holds one:
every field of its manifest has the type the page gives it, every
segment is one of its four kinds, its role is request or response, it
claims at most 256 MiB, every reference points at a record that
landed before it, a part it names is one it has, a copy has the digest
and length it names, and the depth is one more than its base's, or 0,
and at most 32. A manifest's numbers are read as integers the way the
rest of a file's are, with no fraction and no exponent, so 0.5, 0e0 or
a number past a long's range is refused rather than read as another. A
record landed again under an id held before it is not held, as the
Sessionizer's own reader holds only the first; before, a later body
that used a piece only the repeat held made the whole files request
fail.
- Only the row the session holds is served. Any other body record,
request or response, is withheld whole, since what it refers to may be
gone from what is served. A body of another role may be a request, so
it is not held and none of its pieces is, and a response that used one
is not held either. The row held is decided as each record lands, so a
later file cannot change how an earlier one is served.
- The page says that a completion's prompts and LangChain JS's message
id are the OAP's own rules, which the Sessionizer names no record by,
and states the rules a body is held by.
---
docs/en/setup/backend/ai-agent-conversation.md | 25 ++-
.../agent/conversation/providerbody/BodyStore.java | 112 +++++++++--
.../agent/conversation/providerbody/Manifest.java | 129 +++++++++++--
.../agent/conversation/withhold/Withholding.java | 186 +++++++++++-------
.../conversation/providerbody/BodyStoreTest.java | 162 ++++++++++++++++
.../withhold/WithholdingShapesTest.java | 210 ++++++++++++++++++++-
6 files changed, 714 insertions(+), 110 deletions(-)
diff --git a/docs/en/setup/backend/ai-agent-conversation.md b/docs/en/setup/backend/ai-agent-conversation.md
index 33cc87c0f2..7677c88370 100644
--- a/docs/en/setup/backend/ai-agent-conversation.md
+++ b/docs/en/setup/backend/ai-agent-conversation.md
@@ -385,15 +385,30 @@ what is withheld.
completion's `prompts`, becomes `[withheld: system_prompt]`, apart from a block's `type` and `role`, which say what
it is; so does the content of a message of the `system` or `developer` role, wherever it sits, LangChain's pair of
a role and its content, and LangChain's serialized system message, whose role LangChain JS writes only in its
- `id`. Each tool under `tools`, `functions`, `function_declarations`, `tool_definitions` or `available_tools` keeps
- its name, or the name of the one object it wraps, which the transcript's tool calls name anyway, and is otherwise
- `{"withheld":"tool_schemas"}`. The masked request is written
+ `id`. A completion's `prompts` and that `id` are the OAP's own: the Sessionizer names no record by them, since its
+ viewer serves no body. Each tool under `tools`, `functions`, `function_declarations`, `tool_definitions` or
+ `available_tools` keeps its name, or the name of the one object it wraps, which the transcript's tool calls name
+ anyway, and is otherwise `{"withheld":"tool_schemas"}`; a list inside the set is more of its tools, and what the set
+ holds besides them, or is as a whole, is masked like anything else. The masked request is written
back as a body of its own: one `lit` segment, `depth` 0, no `copy` and no `piece`, its own `sha256` and `bytes`,
and `withheld` naming what was applied, so a reader that checks a body against its manifest still can. The request's
own pieces stay except the ones that are a withheld block's text or a tool, whose bytes are the secret, and a body
kept whole in an `unknown` part, which holds them all. A response holds neither and is the stored line, unless it
- refers to a request's bytes, when it is written whole the same way, which hides nothing. A request that does not
- rebuild, or is not one JSON object, is withheld whole, as a transcript record is; so is one whose rebuild would
+ refers to a request's bytes, when it is written whole the same way, which hides nothing; a body that is not UTF-8
+ text cannot be written as a literal byte for byte, and is withheld whole instead. A body is held as the Session
+ Data page says a reader holds one: every field of its manifest has the type the page gives it, every segment is
+ one of its four kinds, a `lit`, a `part`, a `piece` or a `copy`, its role is `request` or `response`, it claims
+ at most 256 MiB, every reference points at a record that landed before it, a part it names is one it has, a copy has the digest and length it names, and the depth is one more
+ than its base's, or 0, and at most 32. A record landed again under an id held before it is not held, as the
+ Sessionizer's own reader holds only the first; the OAP does not check it further, as that reader does, since it
+ is withheld whole whatever it holds. A role that is a string other than those two, or no role at all, is not held
+ either, where the Sessionizer's reader holds it, since such a body may be a request and its pieces may hold the
+ prompt; nor is a segment that is null, sets none of the four, sets more than one or names a piece with no digest,
+ which that reader reads in a way of its own and the page does not give. Only the row
+ the session holds is served: any other body record, request or response, is withheld whole, since what it refers
+ to may be gone from what is served. A request that
+ does not rebuild, or is not one JSON object in UTF-8 text, is withheld whole, as a transcript record is; so is one
+ whose rebuild would
hold more than `maxResponseBytes` at once, counted from the manifests before anything is built: the body, and
beside it the body it copies from, so about two bodies for a chain of copies. So is a record with no manifest of a
schema the OAP knows, since it may be a request;
diff --git a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStore.java b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStore.java
index afee69e8a5..1afcc1e9d5 100644
--- a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStore.java
+++ b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStore.java
@@ -49,15 +49,20 @@ public final class BodyStore {
/** The record and part holding each piece, by the piece's digest; the first record to hold a piece keeps it. */
private final Map<String, int[]> pieceAt = new HashMap<>();
private final Map<String, String> pieceRecord = new HashMap<>();
+ /** Why each record that carries a body is not held, by its id: a size, a reference or a depth the page does not allow. */
+ private final Map<String, String> refused = new HashMap<>();
/** The rebuilt bodies kept, the one kept longest first. */
private final Map<String, byte[]> rebuilt = new LinkedHashMap<>();
private long kept;
private static final class Held {
+ /** The seq of the file the record is in, so the row held is known apart from a later one with its id. */
+ final long seq;
final SessionDataFile.Record record;
final Manifest manifest;
- Held(final SessionDataFile.Record record, final Manifest manifest) {
+ Held(final long seq, final SessionDataFile.Record record, final Manifest manifest) {
+ this.seq = seq;
this.record = record;
this.manifest = manifest;
}
@@ -85,24 +90,42 @@ public final class BodyStore {
*/
public void add(final SessionDataFile file) {
for (final SessionDataFile.Record rec : file.getRecords()) {
- add(rec);
+ add(file.getHeader().getSeq(), rec);
}
}
/**
- * @param rec one record of a <code>provider_body</code> file
- * @return the manifest the record carries, or null when it holds no body
+ * Holds a record's body as the Session Data page's "Provider bodies" section lets a reader: its manifest has the
+ * shape the page gives it, every field of the type it gives and every segment one of its four kinds, its role is
+ * request or response, it claims no more than {@link Manifest#MAX_BYTES}, every reference points at a record that landed
+ * before it, a part it names is one it
+ * has, a copy is of a body with the digest and at least the length it names, and the depth is one more than its
+ * base's, or 0 when it copies nothing, and at most {@link Manifest#MAX_DEPTH}. A record that breaks one is not
+ * held, so whether it is never depends on what lands after it, nor on which later files a reader happens to read.
+ * A record landed again under an id already held, as one is after an interrupted pass, is not held either, as the
+ * Sessionizer's own reader holds only the first: none of its pieces is held, so every piece points at a part of
+ * the record that holds it.
+ *
+ * @param seq the seq of the file the record is in
+ * @param rec one record of a <code>provider_body</code> file, after every record added before it
+ * @return the manifest held under the record's id, or null when it holds no body or is not held
*/
@Nullable
- public Manifest add(final SessionDataFile.Record rec) {
+ private Manifest add(final long seq, final SessionDataFile.Record rec) {
final Manifest m = Manifest.of(rec);
if (m == null || rec.getId() == null || rec.getId().isEmpty()) {
return null;
}
- // a record landed again under its id, as one is after an interrupted pass, is the one already held
- if (!records.containsKey(rec.getId())) {
- records.put(rec.getId(), new Held(rec, m));
+ final Held already = records.get(rec.getId());
+ if (already != null) {
+ return already.manifest;
+ }
+ final String why = refusal(rec.getId(), m);
+ if (why != null) {
+ refused.putIfAbsent(rec.getId(), why);
+ return null;
}
+ records.put(rec.getId(), new Held(seq, rec, m));
final List<SessionDataFile.Part> parts = rec.getParts();
for (int i = 0; i < m.getPartIndex(); i++) {
final SessionDataFile.Part p = parts.get(i);
@@ -119,11 +142,69 @@ public final class BodyStore {
return m;
}
+ /** @return why a record may not be held, as {@link #add(SessionDataFile.Record)} says, or null when it may */
+ @Nullable
+ private String refusal(final String id, final Manifest m) {
+ if (!m.isWellFormed()) {
+ return id + " has a manifest of another shape than the page gives it";
+ }
+ if (!m.isRequest() && !m.isResponse()) {
+ return id + " is a body of neither role the page gives";
+ }
+ if (m.getBytes() < 0 || m.getBytes() > Manifest.MAX_BYTES) {
+ return id + " claims " + m.getBytes() + " bytes, past the " + Manifest.MAX_BYTES + " a body may";
+ }
+ if (m.getDepth() < 0 || m.getDepth() > Manifest.MAX_DEPTH) {
+ return id + " claims depth " + m.getDepth();
+ }
+ int copies = 0;
+ for (final Manifest.Segment seg : m.getSegments()) {
+ if (seg.isUnknown()) {
+ // of the right shape, yet naming nothing a body is made of, such as a piece with no digest
+ return id + " holds a segment of no known kind";
+ }
+ if (seg.getPart() != null && (seg.getPart() < 0 || seg.getPart() >= m.getPartIndex())) {
+ return id + " names part " + seg.getPart() + ", which it does not have";
+ }
+ if (seg.getPiece() != null && !pieceRecord.containsKey(seg.getPiece())) {
+ return id + " refers to piece " + seg.getPiece() + ", which no earlier record holds";
+ }
+ if (seg.getCopyFrom() == null) {
+ continue;
+ }
+ copies++;
+ final Held base = records.get(seg.getCopyFrom());
+ if (base == null) {
+ return id + " copies from " + seg.getCopyFrom() + ", which no earlier record holds";
+ }
+ if (!base.manifest.getSha256().equals(seg.getCopySha256()) || seg.getCopyLen() < 0
+ || seg.getCopyLen() > base.manifest.getBytes()) {
+ return id + " copies from " + seg.getCopyFrom() + " with a digest or length it does not have";
+ }
+ if (m.getDepth() != base.manifest.getDepth() + 1) {
+ return id + " claims depth " + m.getDepth() + " over a base of depth " + base.manifest.getDepth();
+ }
+ }
+ if (copies == 0 && m.getDepth() != 0) {
+ return id + " claims depth " + m.getDepth() + " and copies nothing";
+ }
+ return null;
+ }
+
/** @return the bytes of the rebuilt bodies kept */
long keptBytes() {
return kept;
}
+ /**
+ * @return whether the record on this row of the file with this seq is the one held under its id: not one the page
+ * does not let a reader hold, and not one landed again under an id held before it
+ */
+ public boolean holds(final long seq, final int row, final String id) {
+ final Held h = records.get(id);
+ return h != null && h.seq == seq && h.record.getRow() == row;
+ }
+
/**
* @return the manifest held under the record id, or null
*/
@@ -147,7 +228,7 @@ public final class BodyStore {
}
final Held h = records.get(id);
if (h == null) {
- throw new BodyException("the session holds no body " + id);
+ throw new BodyException(refused.getOrDefault(id, "the session holds no body " + id));
}
// Whether a body fits is decided from the manifests, before anything is built, so it depends on the body
// and the bodies it copies from alone, never on which bodies this store happens to keep: a file asked for
@@ -187,9 +268,10 @@ public final class BodyStore {
built += Math.max(0, seg.getCopyLen());
} else if (seg.getLit() != null) {
built += seg.getLit().getBytes(StandardCharsets.UTF_8).length;
- } else if (seg.getPart() >= 0 && seg.getPart() < m.getPartIndex()) {
- built += Math.max(0, h.record.getParts().get(seg.getPart()).getBytes());
+ } else if (seg.getPart() != null && seg.getPart() >= 0 && seg.getPart() < m.getPartIndex()) {
+ built += Math.max(0, h.record.getParts().get(seg.getPart().intValue()).getBytes());
} else if (seg.getPiece() != null && pieceRecord.containsKey(seg.getPiece())) {
+ // a piece is held only by the record whose part it is, so its index is within that record's parts
final Held holder = records.get(pieceRecord.get(seg.getPiece()));
built += Math.max(0, holder.record.getParts().get(pieceAt.get(seg.getPiece())[0]).getBytes());
}
@@ -252,7 +334,7 @@ public final class BodyStore {
final byte[] bytes;
if (seg.getLit() != null) {
bytes = seg.getLit().getBytes(StandardCharsets.UTF_8);
- } else if (seg.getPart() >= 0) {
+ } else if (seg.getPart() != null) {
bytes = partBytes(h, seg.getPart());
} else if (seg.getPiece() != null) {
bytes = pieceBytes(id, seg.getPiece());
@@ -273,12 +355,12 @@ public final class BodyStore {
return body;
}
- private static byte[] partBytes(final Held h, final int index) throws BodyException {
+ private static byte[] partBytes(final Held h, final long index) throws BodyException {
final List<SessionDataFile.Part> parts = h.record.getParts();
- if (index >= h.manifest.getPartIndex()) {
+ if (index < 0 || index >= h.manifest.getPartIndex()) {
throw new BodyException(h.record.getId() + " names part " + index + ", which it does not have");
}
- final SessionDataFile.Part part = parts.get(index);
+ final SessionDataFile.Part part = parts.get((int) index);
final String data = part.data();
if (data == null) {
throw new BodyException(h.record.getId() + " names part " + index + ", which holds no data");
diff --git a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/Manifest.java b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/Manifest.java
index f1b7c93ad1..a131aee926 100644
--- a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/Manifest.java
+++ b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/Manifest.java
@@ -21,6 +21,7 @@ import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import javax.annotation.Nullable;
@@ -46,16 +47,23 @@ public final class Manifest {
public static final long MAX_BYTES = 256L * 1024 * 1024;
/** How many copies may lie between a body and one with none, as the page says. */
public static final int MAX_DEPTH = 32;
+ /** The four kinds of segment the page gives, each a key of its own. */
+ private static final List<String> SEGMENT_KINDS = Arrays.asList("lit", "part", "piece", "copy");
+ /** The manifest's fields the page gives as text. */
+ private static final List<String> STRING_KEYS = Arrays.asList(
+ "schema", "role", "src", "sha256", "chain", "why", "model", "session", "run", "call", "request", "previous_request");
/** The manifest as written, so a rewrite keeps every field it does not change. */
private final JsonObject json;
private final String role;
private final String sha256;
private final long bytes;
- private final int depth;
+ private final long depth;
private final List<Segment> segments;
/** The part the manifest is, so the parts before it are the record's own pieces. */
private final int partIndex;
+ /** Whether the manifest has the shape the page gives it; see {@link #wellFormed(JsonObject)}. */
+ private final boolean wellFormed;
private Manifest(final JsonObject json, final int partIndex) {
this.json = json;
@@ -63,7 +71,8 @@ public final class Manifest {
this.role = string(json, "role");
this.sha256 = string(json, "sha256");
this.bytes = longOf(json, "bytes");
- this.depth = (int) longOf(json, "depth");
+ this.depth = longOf(json, "depth");
+ this.wellFormed = wellFormed(json);
final List<Segment> list = new ArrayList<>();
final JsonElement segs = json.get("segments");
if (segs != null && segs.isJsonArray()) {
@@ -101,6 +110,10 @@ public final class Manifest {
return ROLE_REQUEST.equals(role);
}
+ public boolean isResponse() {
+ return ROLE_RESPONSE.equals(role);
+ }
+
/**
* @return the <code>why</code> the page sets when the body was kept whole in one part because it could not
* be cut, or null
@@ -116,15 +129,100 @@ public final class Manifest {
return e != null && e.isJsonPrimitive() && e.getAsJsonPrimitive().isString() ? e.getAsString() : "";
}
+ /**
+ * Whether the manifest has the shape the page gives it: every field it gives has the type it gives it, where it is
+ * present and not null, its text fields text, its numbers integers, and its segments a list in which each is one of
+ * the four kinds the page gives, a <code>lit</code>, a <code>part</code>, a <code>piece</code> or a
+ * <code>copy</code>, with that one key set to a value of its type, and a copy's own fields theirs. A segment that is
+ * null, sets none of the four or sets more than one is none of them: the page says nothing of what it adds to a
+ * body, so it is refused rather than read one way or another. The Sessionizer's reader, decoding into its own
+ * types, reads a null or empty segment as one that adds nothing.
+ */
+ private static boolean wellFormed(final JsonObject json) {
+ for (final String key : STRING_KEYS) {
+ if (!textOrAbsent(json.get(key))) {
+ return false;
+ }
+ }
+ if (!integerOrAbsent(json.get("bytes")) || !integerOrAbsent(json.get("depth"))) {
+ return false;
+ }
+ final JsonElement segs = json.get("segments");
+ if (segs == null || segs.isJsonNull()) {
+ return true;
+ }
+ if (!segs.isJsonArray()) {
+ return false;
+ }
+ for (final JsonElement e : segs.getAsJsonArray()) {
+ if (!e.isJsonObject()) {
+ return false;
+ }
+ final JsonObject seg = e.getAsJsonObject();
+ if (!textOrAbsent(seg.get("lit")) || !textOrAbsent(seg.get("piece")) || !integerOrAbsent(seg.get("part"))) {
+ return false;
+ }
+ int kinds = 0;
+ for (final String kind : SEGMENT_KINDS) {
+ if (seg.get(kind) != null && !seg.get(kind).isJsonNull()) {
+ kinds++;
+ }
+ }
+ if (kinds != 1) {
+ return false;
+ }
+ final JsonElement copy = seg.get("copy");
+ if (copy == null || copy.isJsonNull()) {
+ continue;
+ }
+ if (!copy.isJsonObject()) {
+ return false;
+ }
+ final JsonObject c = copy.getAsJsonObject();
+ if (!textOrAbsent(c.get("from")) || !textOrAbsent(c.get("sha256")) || !integerOrAbsent(c.get("len"))) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private static boolean textOrAbsent(@Nullable final JsonElement e) {
+ return e == null || e.isJsonNull() || e.isJsonPrimitive() && e.getAsJsonPrimitive().isString();
+ }
+
+ private static boolean integerOrAbsent(@Nullable final JsonElement e) {
+ if (e == null || e.isJsonNull()) {
+ return true;
+ }
+ if (!e.isJsonPrimitive() || !e.getAsJsonPrimitive().isNumber()) {
+ return false;
+ }
+ try {
+ Long.parseLong(e.getAsString());
+ return true;
+ } catch (final NumberFormatException ex) {
+ return false;
+ }
+ }
+
+ /**
+ * @return the integer under the key: 0 when the key is absent, and -1, which no field of a manifest may hold,
+ * when the value is not an integer as the page writes one, with no fraction and no exponent, in the range of a
+ * long, so a value that is not is refused rather than read as another, as <code>0.5</code>, <code>0e0</code> or a
+ * number past the range would be. It reads a number as the other integer fields of a file are read.
+ */
static long longOf(final JsonObject json, final String key) {
final JsonElement e = json.get(key);
- if (e == null || !e.isJsonPrimitive() || !e.getAsJsonPrimitive().isNumber()) {
+ if (e == null || e.isJsonNull()) {
return 0;
}
+ if (!e.isJsonPrimitive() || !e.getAsJsonPrimitive().isNumber()) {
+ return -1;
+ }
try {
- return e.getAsLong();
+ return Long.parseLong(e.getAsString());
} catch (final NumberFormatException ex) {
- return 0;
+ return -1;
}
}
@@ -136,8 +234,9 @@ public final class Manifest {
/** Literal text whose UTF-8 bytes are bytes of the body. */
@Nullable
private final String lit;
- /** The index of a part of this record, or -1. */
- private final int part;
+ /** The index of a part of this record, or null when the segment is not a part; -1 when it is not an index. */
+ @Nullable
+ private final Long part;
/** The digest of a piece an earlier record holds, or null. */
@Nullable
private final String piece;
@@ -152,7 +251,7 @@ public final class Manifest {
/** Set when the segment is none of the four, so a rebuild can refuse it. */
private final boolean unknown;
- private Segment(@Nullable final String lit, final int part, @Nullable final String piece,
+ private Segment(@Nullable final String lit, @Nullable final Long part, @Nullable final String piece,
@Nullable final String copyFrom, @Nullable final String copySha256, final long copyLen,
final boolean unknown) {
this.lit = lit;
@@ -166,27 +265,27 @@ public final class Manifest {
static Segment of(final JsonElement e) {
if (!e.isJsonObject()) {
- return new Segment(null, -1, null, null, null, 0, true);
+ return new Segment(null, null, null, null, null, 0, true);
}
final JsonObject o = e.getAsJsonObject();
final JsonElement lit = o.get("lit");
if (lit != null && lit.isJsonPrimitive() && lit.getAsJsonPrimitive().isString()) {
- return new Segment(lit.getAsString(), -1, null, null, null, 0, false);
+ return new Segment(lit.getAsString(), null, null, null, null, 0, false);
}
final JsonElement part = o.get("part");
- if (part != null && part.isJsonPrimitive() && part.getAsJsonPrimitive().isNumber()) {
- return new Segment(null, (int) longOf(o, "part"), null, null, null, 0, false);
+ if (part != null && !part.isJsonNull()) {
+ return new Segment(null, longOf(o, "part"), null, null, null, 0, false);
}
final String piece = string(o, "piece");
if (!piece.isEmpty()) {
- return new Segment(null, -1, piece, null, null, 0, false);
+ return new Segment(null, null, piece, null, null, 0, false);
}
final JsonElement copy = o.get("copy");
if (copy != null && copy.isJsonObject()) {
final JsonObject c = copy.getAsJsonObject();
- return new Segment(null, -1, null, string(c, "from"), string(c, "sha256"), longOf(c, "len"), false);
+ return new Segment(null, null, null, string(c, "from"), string(c, "sha256"), longOf(c, "len"), false);
}
- return new Segment(null, -1, null, null, null, 0, true);
+ return new Segment(null, null, null, null, null, 0, true);
}
/**
diff --git a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/Withholding.java b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/Withholding.java
index 57234bdc3b..2c65e941b6 100644
--- a/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/Withholding.java
+++ b/oap-server/analyzer/ai-agent-conversation/src/main/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/Withholding.java
@@ -24,6 +24,9 @@ import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonPrimitive;
import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import java.nio.charset.CharacterCodingException;
+import java.nio.charset.CodingErrorAction;
import java.nio.charset.StandardCharsets;
import java.util.ArrayDeque;
import java.util.ArrayList;
@@ -56,7 +59,8 @@ import org.apache.skywalking.oap.server.library.util.StringUtil;
* is withheld whole rather than served with a reference that leads nowhere. A response holds neither name and is
* served as stored, unless it refers to a request's bytes, when it is written whole the same way. A record of a
* <code>provider_body</code> file with no manifest this reader knows is withheld whole, since nothing says what it
- * holds. A file with nothing to withhold is served byte for byte.
+ * holds, and so is any row that is not the one its session holds, as {@link BodyStore#holds} says. A file with
+ * nothing to withhold is served byte for byte.
*
* <p>A rewritten file gets a closing line of its own, so it still reads and checks on its own. Its records keep
* their rows, so every reference into it still resolves. A file whose stored closing line does not check keeps
@@ -143,7 +147,7 @@ public final class Withholding {
for (final SessionDataFile.Record rec : file.getRecords()) {
String line = null;
if (bodies) {
- line = body(rec, store, hidden, rewritten);
+ line = body(header.getSeq(), rec, store, hidden, rewritten);
} else if (Hide.carriesAny(rec.flags(), hidden)) {
line = withheldLine(rec);
}
@@ -214,6 +218,23 @@ public final class Withholding {
return -1;
}
+ /**
+ * @return the bytes as text, or null when they are not UTF-8: replacing what does not decode would change the
+ * bytes a digest was taken of
+ */
+ @Nullable
+ private static String utf8(final byte[] bytes) {
+ try {
+ return StandardCharsets.UTF_8.newDecoder()
+ .onMalformedInput(CodingErrorAction.REPORT)
+ .onUnmappableCharacter(CodingErrorAction.REPORT)
+ .decode(ByteBuffer.wrap(bytes))
+ .toString();
+ } catch (final CharacterCodingException e) {
+ return null;
+ }
+ }
+
private static void write(final ByteArrayOutputStream out, final String line) {
final byte[] bytes = line.getBytes(StandardCharsets.UTF_8);
out.write(bytes, 0, bytes.length);
@@ -224,20 +245,23 @@ public final class Withholding {
* @return the served line of one record of a <code>provider_body</code> file, or null for the stored line: a
* request is masked, or withheld whole when it cannot be rebuilt or has no id to be held under; a response is
* written whole when it refers to a request's bytes, and withheld whole when it cannot be; a record with no
- * manifest this reader knows, such as one of a later schema, is withheld whole, since it may be a request
+ * manifest this reader knows, such as one of a later schema, is withheld whole, since it may be a request, and so
+ * is a row that is not the one the session holds, a role other than request or response among the reasons
*/
@Nullable
- private static String body(final SessionDataFile.Record rec, @Nullable final BodyStore store, final List<String> hidden,
- final Map<String, Boolean> rewritten) {
+ private static String body(final long seq, final SessionDataFile.Record rec, @Nullable final BodyStore store,
+ final List<String> hidden, final Map<String, Boolean> rewritten) {
final Manifest m = Manifest.of(rec);
- if (m == null) {
+ // Only the row the session holds is served. One with no id, or a role, a size, a reference or a depth the page
+ // does not allow, or one landed again under an id held before it, is withheld whole: a reader holds the first
+ // and refuses or passes over the rest, and what such a row refers to may be gone from what is served.
+ if (m == null || store == null || StringUtil.isEmpty(rec.getId()) || !store.holds(seq, rec.getRow(), rec.getId())) {
return withheldLine(rec);
}
- final Manifest held = store == null || StringUtil.isEmpty(rec.getId()) ? null : store.manifest(rec.getId());
if (m.isRequest()) {
- return held == null ? withheldLine(rec) : request(rec, m, store, hidden, rewritten);
+ return request(rec, m, store, hidden, rewritten);
}
- if (held != null && leansOnRewritten(rec.getId(), m, store, rewritten)) {
+ if (leansOnRewritten(rec.getId(), m, store, rewritten)) {
final String whole = standalone(rec, m, store);
return whole == null ? withheldLine(rec) : whole;
}
@@ -327,20 +351,24 @@ public final class Withholding {
/**
* @return the record's line with its body rebuilt and written as one literal segment, its own pieces kept, and
- * its digest and size as they were; or null when the body does not rebuild
+ * its digest and size as they were; or null when the body does not rebuild, or is not UTF-8 text, which a literal
+ * segment cannot hold byte for byte
*/
@Nullable
private static String standalone(final SessionDataFile.Record rec, final Manifest m, final BodyStore store) {
- final byte[] body;
+ final String body;
try {
- body = store.body(rec.getId());
+ body = utf8(store.body(rec.getId()));
} catch (final BodyStore.BodyException e) {
return null;
}
+ if (body == null) {
+ return null;
+ }
final JsonObject manifest = m.getJson().deepCopy();
manifest.addProperty("depth", 0);
manifest.remove("segments");
- manifest.add("segments", Manifest.Segment.literal(new String(body, StandardCharsets.UTF_8)));
+ manifest.add("segments", Manifest.Segment.literal(body));
final List<String> parts = new ArrayList<>();
final List<SessionDataFile.Part> own = rec.getParts();
for (int i = 0; i < m.getPartIndex(); i++) {
@@ -361,22 +389,24 @@ public final class Withholding {
/**
* @return the request record's line with its body masked; the line written whole when the body holds nothing
* to withhold but refers to a request's bytes; null for the stored line when it holds nothing and refers to
- * none; the record withheld whole when its body does not rebuild or is not one JSON object
+ * none; the record withheld whole when its body does not rebuild or is not one JSON object in UTF-8 text
*/
@Nullable
private static String request(final SessionDataFile.Record rec, final Manifest m, final BodyStore store,
final List<String> hidden, final Map<String, Boolean> rewritten) {
- final JsonObject body;
+ final String text;
try {
- final JsonElement parsed = Schema.parse(new String(store.body(rec.getId()), StandardCharsets.UTF_8));
- body = parsed != null && parsed.isJsonObject() ? parsed.getAsJsonObject() : null;
+ text = utf8(store.body(rec.getId()));
} catch (final BodyStore.BodyException e) {
return withheldLine(rec);
}
- if (body == null) {
- // a body that is not one JSON object cannot be read for what it holds, so all of it is withheld
+ final JsonElement parsed = text == null ? null : Schema.parse(text);
+ if (parsed == null || !parsed.isJsonObject()) {
+ // a body that is not one JSON object in UTF-8 text cannot be read for what it holds, so all of it is
+ // withheld
return withheldLine(rec);
}
+ final JsonObject body = parsed.getAsJsonObject();
// every string of what is masked, so a piece that held one goes with it
final Set<String> maskedTexts = new HashSet<>();
final boolean changed = mask(body, hidden, maskedTexts);
@@ -467,44 +497,61 @@ public final class Withholding {
* @return whether anything was masked
*/
private static boolean mask(final JsonElement v, final List<String> hidden, final Set<String> maskedTexts) {
- final boolean prompt = hidden.contains(Hide.SYSTEM_PROMPT);
- boolean changed = false;
+ boolean changed = maskShape(v, hidden, maskedTexts);
if (v.isJsonObject()) {
final JsonObject o = v.getAsJsonObject();
- if (prompt) {
- changed = maskSystemMessage(o, maskedTexts);
- }
for (final String key : new ArrayList<>(o.keySet())) {
- final JsonElement inner = o.get(key);
- final String k = keyOf(key);
- if (prompt && PROMPT_KEYS.contains(k) && holdsSomething(inner)) {
- o.add(key, maskedPrompt(inner, maskedTexts));
- changed = true;
- } else if (hidden.contains(Hide.TOOL_SCHEMAS) && TOOL_KEYS.contains(k) && offersTools(inner)) {
- changed |= maskTools(o, key, maskedTexts);
- } else {
- changed |= mask(inner, hidden, maskedTexts);
- }
+ changed |= maskMember(o, key, hidden, maskedTexts);
}
} else if (v.isJsonArray()) {
- final JsonArray list = v.getAsJsonArray();
- if (prompt && systemPairIn(list)) {
- for (final JsonElement item : list) {
- final JsonArray pair = item.isJsonArray() ? item.getAsJsonArray() : null;
- if (pair != null && systemRole(pair.get(0)) && holdsSomething(pair.get(1))) {
- collect(pair.get(1), maskedTexts);
- pair.set(1, new JsonPrimitive(MARKER_SYSTEM));
- changed = true;
- }
- }
- }
- for (final JsonElement item : list) {
+ for (final JsonElement item : v.getAsJsonArray()) {
changed |= mask(item, hidden, maskedTexts);
}
}
return changed;
}
+ /**
+ * Masks what a value is as a whole, before its members are: a system message, when it is an object, and the system
+ * messages of a list of messages written as LangChain's pairs, when it is a list. Every value the walk reaches is
+ * looked at this way, a set of tools too.
+ */
+ private static boolean maskShape(final JsonElement v, final List<String> hidden, final Set<String> maskedTexts) {
+ if (!hidden.contains(Hide.SYSTEM_PROMPT)) {
+ return false;
+ }
+ if (v.isJsonObject()) {
+ return maskSystemMessage(v.getAsJsonObject(), maskedTexts);
+ }
+ if (!v.isJsonArray() || !systemPairIn(v.getAsJsonArray())) {
+ return false;
+ }
+ boolean changed = false;
+ for (final JsonElement item : v.getAsJsonArray()) {
+ final JsonArray pair = item.isJsonArray() ? item.getAsJsonArray() : null;
+ if (pair != null && systemRole(pair.get(0)) && holdsSomething(pair.get(1))) {
+ collect(pair.get(1), maskedTexts);
+ pair.set(1, new JsonPrimitive(MARKER_SYSTEM));
+ changed = true;
+ }
+ }
+ return changed;
+ }
+
+ /** Masks one member of an object: a prompt under a prompt key, tools under a tool key, and otherwise what it holds. */
+ private static boolean maskMember(final JsonObject o, final String key, final List<String> hidden, final Set<String> maskedTexts) {
+ final JsonElement inner = o.get(key);
+ final String k = keyOf(key);
+ if (hidden.contains(Hide.SYSTEM_PROMPT) && PROMPT_KEYS.contains(k) && holdsSomething(inner)) {
+ o.add(key, maskedPrompt(inner, maskedTexts));
+ return true;
+ }
+ if (hidden.contains(Hide.TOOL_SCHEMAS) && TOOL_KEYS.contains(k) && offersTools(inner)) {
+ return maskTools(o, key, hidden, maskedTexts);
+ }
+ return mask(inner, hidden, maskedTexts);
+ }
+
/** A key as the prompt and tool keys are compared: without case, underscores or dashes. */
private static String keyOf(final String key) {
return key.toLowerCase(Locale.ROOT).replace("_", "").replace("-", "");
@@ -648,38 +695,47 @@ public final class Withholding {
message.addProperty("content", MARKER_SYSTEM);
}
+ /** The set of tools under a tool key, masked as {@link #maskToolSet} says. */
+ private static boolean maskTools(final JsonObject holder, final String key, final List<String> hidden,
+ final Set<String> maskedTexts) {
+ final JsonElement tools = holder.get(key);
+ return tools != null && maskToolSet(tools, hidden, maskedTexts);
+ }
+
/**
- * One set of tools, a list or an object keyed by name: each keeps its name, since the transcript's tool calls
- * name it anyway, and loses the rest. Every string of a tool is remembered, so a piece that held one, such as a
- * long description, goes with it.
+ * One set of tools, a list or an object keyed by name. What the set is as a whole is masked first, as every value
+ * is. Then each tool keeps its name, since the transcript's tool calls name it anyway, and loses the rest, and
+ * every string of it is remembered, so a piece that held one, such as a long description, goes with it. A list
+ * inside the set is more of its tools. Anything else in it is masked like any other member, as the Sessionizer's
+ * walk finds it. A masked tool is not walked again, so a tool named like a prompt key keeps the tools' marker.
*/
- private static boolean maskTools(final JsonObject holder, final String key, final Set<String> maskedTexts) {
- final JsonElement tools = holder.get(key);
- if (tools == null) {
- return false;
- }
- boolean changed = false;
+ private static boolean maskToolSet(final JsonElement tools, final List<String> hidden, final Set<String> maskedTexts) {
+ boolean changed = maskShape(tools, hidden, maskedTexts);
if (tools.isJsonArray()) {
final JsonArray list = tools.getAsJsonArray();
for (int i = 0; i < list.size(); i++) {
final JsonElement tool = list.get(i);
- if (!tool.isJsonObject()) {
- continue;
+ if (tool.isJsonObject()) {
+ collect(tool, maskedTexts);
+ list.set(i, maskedTool(nameOf(tool.getAsJsonObject())));
+ changed = true;
+ } else if (tool.isJsonArray()) {
+ changed |= maskToolSet(tool, hidden, maskedTexts);
}
- collect(tool, maskedTexts);
- list.set(i, maskedTool(nameOf(tool.getAsJsonObject())));
- changed = true;
}
} else if (tools.isJsonObject()) {
final JsonObject byName = tools.getAsJsonObject();
for (final String name : new ArrayList<>(byName.keySet())) {
final JsonElement tool = byName.get(name);
- if (!tool.isJsonObject()) {
- continue;
+ if (tool.isJsonObject()) {
+ collect(tool, maskedTexts);
+ byName.add(name, maskedTool(null));
+ changed = true;
+ } else if (tool.isJsonArray() && !PROMPT_KEYS.contains(keyOf(name))) {
+ changed |= maskToolSet(tool, hidden, maskedTexts);
+ } else {
+ changed |= maskMember(byName, name, hidden, maskedTexts);
}
- collect(tool, maskedTexts);
- byName.add(name, maskedTool(null));
- changed = true;
}
}
return changed;
diff --git a/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStoreTest.java b/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStoreTest.java
index 0b19635456..4e46b0153e 100644
--- a/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStoreTest.java
+++ b/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/providerbody/BodyStoreTest.java
@@ -18,6 +18,9 @@
package org.apache.skywalking.oap.server.ai.agent.conversation.providerbody;
import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
import java.util.Map;
import org.apache.skywalking.oap.server.ai.agent.conversation.Fixtures;
import org.apache.skywalking.oap.server.ai.agent.conversation.format.Digests;
@@ -29,6 +32,7 @@ import org.junit.jupiter.params.provider.ValueSource;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -150,6 +154,164 @@ public class BodyStoreTest {
+ "\",\"bytes\":" + body.length() + ",\"depth\":" + depth + ",\"segments\":" + segments + "}}]}\n";
}
+ /**
+ * A record landed again under its id is the one already held, and none of its pieces is held again: a later body
+ * that uses a piece only the repeat holds is not held either, and says why, rather than reading a part the held
+ * record does not have.
+ */
+ @Test
+ public void aRecordLandedAgainIsTheOneHeld() throws Exception {
+ final String p0 = "\"" + "p".repeat(20) + "\"";
+ final String p1 = "\"" + "q".repeat(20) + "\"";
+ final String later = "{\"b\":" + p1 + "}";
+ final String file = HEADER
+ + record("x", "{\"a\":1}", 0, "[{\"lit\":\"{\\\"a\\\":1}\"}]")
+ + withPieces("x", "[" + p0 + "," + p1 + "]", Arrays.asList(p0, p1), "[{\"lit\":\"[\"},{\"part\":0},{\"lit\":\",\"},{\"part\":1},{\"lit\":\"]\"}]")
+ + record("y", later, 0, "[{\"lit\":\"{\\\"b\\\":\"},{\"piece\":\"" + Digests.sha256Hex(p1.getBytes(StandardCharsets.UTF_8)) + "\"},{\"lit\":\"}\"}]")
+ + "{\"t\":\"end\",\"records\":3,\"digest\":\"x\"}\n";
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(file.getBytes(StandardCharsets.UTF_8)));
+ assertEquals("{\"a\":1}", new String(store.body("x"), StandardCharsets.UTF_8));
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> store.body("y")).getMessage().contains("which no earlier record holds"));
+ }
+
+ /**
+ * A reference to a record that lands later is not held, as the page says no reference may be, so a body reads the
+ * same whether or not the later file is read: here it is not held either way.
+ */
+ @Test
+ public void aReferenceToALaterRecordIsNotHeldWhateverIsRead() throws Exception {
+ final String p = "\"" + "r".repeat(20) + "\"";
+ final String early = HEADER
+ + record("e", "{\"e\":" + p + "}", 0, "[{\"lit\":\"{\\\"e\\\":\"},{\"piece\":\"" + Digests.sha256Hex(p.getBytes(StandardCharsets.UTF_8)) + "\"},{\"lit\":\"}\"}]")
+ + "{\"t\":\"end\",\"records\":1,\"digest\":\"x\"}\n";
+ final String late = HEADER.replace("\"seq\":1", "\"seq\":2")
+ + withPieces("l", "[" + p + "]", Collections.singletonList(p), "[{\"lit\":\"[\"},{\"part\":0},{\"lit\":\"]\"}]")
+ + "{\"t\":\"end\",\"records\":1,\"digest\":\"x\"}\n";
+ final BodyStore alone = new BodyStore();
+ alone.add(SessionDataFile.parse(early.getBytes(StandardCharsets.UTF_8)));
+ final BodyStore both = new BodyStore();
+ both.add(SessionDataFile.parse(early.getBytes(StandardCharsets.UTF_8)));
+ both.add(SessionDataFile.parse(late.getBytes(StandardCharsets.UTF_8)));
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> alone.body("e")).getMessage().contains("which no earlier record holds"));
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> both.body("e")).getMessage().contains("which no earlier record holds"));
+ }
+
+ /** A depth the page does not allow is not held: one claimed without a copy, and one that is not its base's plus one. */
+ @Test
+ public void aDepthThePageDoesNotAllowIsNotHeld() throws Exception {
+ final String file = HEADER
+ + record("a", "a".repeat(20), 3, "[{\"lit\":\"" + "a".repeat(20) + "\"}]")
+ + record("b", "b".repeat(20), 0, "[{\"lit\":\"" + "b".repeat(20) + "\"}]")
+ + record("c", "b".repeat(10), 2, "[{\"copy\":{\"from\":\"b\",\"sha256\":\"" + Digests.sha256Hex("b".repeat(20).getBytes(StandardCharsets.UTF_8))
+ + "\",\"len\":10}}]")
+ + "{\"t\":\"end\",\"records\":3,\"digest\":\"x\"}\n";
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(file.getBytes(StandardCharsets.UTF_8)));
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> store.body("a")).getMessage().contains("claims depth 3 and copies nothing"));
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> store.body("c")).getMessage().contains("claims depth 2 over a base of depth 0"));
+ assertEquals(20, store.body("b").length);
+ }
+
+ /**
+ * What the page does not let a reader hold is not held, whatever a later rebuild would say: a part the record does
+ * not have, or a negative one; a depth past 32 over a base at 32; a size past 256 MiB, so a later record that uses
+ * its piece is not held either; a depth past what an int holds, which must not wrap to a small one; and an integer
+ * spelled with a fraction or an exponent, or past what a long holds, which must not be read as another.
+ */
+ @Test
+ public void whatThePageDoesNotAllowIsNotHeld() {
+ final String p = "\"" + "s".repeat(20) + "\"";
+ final String deep = chain(33);
+ final StringBuilder file = new StringBuilder(deep.substring(0, deep.lastIndexOf("{\"t\":\"end\"")));
+ final String base = chainBody(32);
+ file.append(record("b33", base.substring(0, 10), 33, "[{\"copy\":{\"from\":\"b32\",\"sha256\":\""
+ + Digests.sha256Hex(base.getBytes(StandardCharsets.UTF_8)) + "\",\"len\":10}}]"));
+ file.append(withPieces("past", "[" + p + "," + p + "]", Collections.singletonList(p), "[{\"lit\":\"[\"},{\"part\":1},{\"lit\":\"]\"}]"));
+ file.append(withPieces("minus", "[" + p + "]", Collections.singletonList(p), "[{\"lit\":\"[\"},{\"part\":-1},{\"lit\":\"]\"}]"));
+ file.append(withPieces("big", "[" + p + "]", Collections.singletonList(p), "[{\"lit\":\"[\"},{\"part\":0},{\"lit\":\"]\"}]")
+ .replace("\"bytes\":" + ("[" + p + "]").length() + ",", "\"bytes\":300000000,"));
+ file.append(record("user", "[" + p + "]", 0, "[{\"lit\":\"[\"},{\"piece\":\"" + Digests.sha256Hex(p.getBytes(StandardCharsets.UTF_8)) + "\"},{\"lit\":\"]\"}]"));
+ file.append(record("wrap", "w".repeat(20), 0, "[{\"lit\":\"" + "w".repeat(20) + "\"}]").replace("\"depth\":0", "\"depth\":4294967296"));
+ // an integer is written with no fraction and no exponent, as the page writes one, and a long holds it
+ final String[] spelled = {"\"depth\":0.5", "\"depth\":0.0", "\"depth\":0e0", "\"depth\":18446744073709551616", "\"bytes\":20.0,\"depth\":0"};
+ for (int i = 0; i < spelled.length; i++) {
+ final String rec = record("spelled" + i, "z".repeat(20), 0, "[{\"lit\":\"" + "z".repeat(20) + "\"}]");
+ file.append(spelled[i].startsWith("\"bytes") ? rec.replace("\"bytes\":20,\"depth\":0", spelled[i]) : rec.replace("\"depth\":0", spelled[i]));
+ }
+ file.append("{\"t\":\"end\",\"records\":44,\"digest\":\"x\"}\n");
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(file.toString().getBytes(StandardCharsets.UTF_8)));
+ assertNotNull(store.manifest("b32"), "the base at depth 32 is held");
+ for (final String id : new String[] {"b33", "past", "minus", "big", "user", "wrap", "spelled0", "spelled1", "spelled2", "spelled3", "spelled4"}) {
+ assertNull(store.manifest(id), id + " is held");
+ }
+ assertTrue(assertThrows(BodyStore.BodyException.class, () -> store.body("user")).getMessage().contains("which no earlier record holds"));
+ }
+
+ /**
+ * A manifest of another shape than the page gives it is not held, and neither is a body that uses its piece: a
+ * field of another type, as a reader that decodes it as the page writes it refuses it, and a segment that is not
+ * one of the page's four kinds, null, empty, setting two or naming a piece with no digest, which the page says
+ * nothing of.
+ */
+ @Test
+ public void aManifestOfAnotherShapeIsNotHeld() {
+ final String body = "t".repeat(20);
+ final String lit = "[{\"lit\":\"" + body + "\"}]";
+ final String sha = "\"sha256\":\"" + Digests.sha256Hex(body.getBytes(StandardCharsets.UTF_8)) + "\"";
+ final String p = "\"" + "u".repeat(20) + "\"";
+ final String[] typed = {
+ record("digest", body, 0, lit).replace(sha, "\"sha256\":5"),
+ record("src", body, 0, lit).replace("\"role\":\"request\"", "\"role\":\"request\",\"src\":true"),
+ record("segment", body, 0, "[5]"),
+ record("text", body, 0, "[{\"lit\":5}]"),
+ record("piece", body, 0, "[{\"piece\":5}]"),
+ record("copy", body, 0, "[{\"copy\":{\"from\":5,\"sha256\":\"x\",\"len\":1}}]"),
+ record("segments", body, 0, "{}"),
+ withPieces("holder", "[" + p + "]", Collections.singletonList(p), "[{\"lit\":\"[\"},{\"part\":0},{\"lit\":\"]\"}]")
+ .replace("\"role\":\"request\"", "\"role\":\"request\",\"model\":[]"),
+ record("user", "[" + p + "]", 0, "[{\"lit\":\"[\"},{\"piece\":\"" + Digests.sha256Hex(p.getBytes(StandardCharsets.UTF_8)) + "\"},{\"lit\":\"]\"}]"),
+ record("empty", body, 0, "[{\"lit\":\"" + body + "\"},{}]"),
+ record("null", body, 0, "[{\"lit\":\"" + body + "\"},null]"),
+ record("two", body, 0, "[{\"lit\":\"" + body + "\",\"part\":7}]"),
+ record("nodigest", body, 0, "[{\"lit\":\"" + body + "\"},{\"piece\":\"\"}]"),
+ record("plain", body, 0, lit),
+ };
+ final StringBuilder file = new StringBuilder(HEADER);
+ for (final String rec : typed) {
+ file.append(rec);
+ }
+ file.append("{\"t\":\"end\",\"records\":").append(typed.length).append(",\"digest\":\"x\"}\n");
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(file.toString().getBytes(StandardCharsets.UTF_8)));
+ for (final String id : new String[] {"digest", "src", "segment", "text", "piece", "copy", "segments", "holder", "user", "empty", "null", "two", "nodigest"}) {
+ assertNull(store.manifest(id), id + " is held");
+ }
+ assertNotNull(store.manifest("plain"), "a manifest of the page's shape is held");
+ }
+
+ /** The body of b(i) in {@link #chain(int)}. */
+ private static String chainBody(final int i) {
+ String body = null;
+ for (int j = 0; j <= i; j++) {
+ final String own = String.valueOf((char) ('a' + j % 26)).repeat(j == 0 ? 60 : 50);
+ body = body == null ? own : body.substring(0, 10) + own;
+ }
+ return body;
+ }
+
+ /** A request record with its pieces as data parts before the manifest. */
+ private static String withPieces(final String id, final String body, final List<String> pieces, final String segments) {
+ final StringBuilder parts = new StringBuilder();
+ for (final String piece : pieces) {
+ parts.append("{\"k\":\"data\",\"data\":").append(piece).append(",\"state\":\"available\",\"bytes\":").append(piece.length()).append("},");
+ }
+ return "{\"ord\":1,\"off\":0,\"sha\":\"0\",\"bytes\":1,\"id\":\"" + id + "\",\"parts\":[" + parts
+ + "{\"k\":\"data\",\"data\":{\"schema\":\"provider_body/1\",\"role\":\"request\",\"sha256\":\""
+ + Digests.sha256Hex(body.getBytes(StandardCharsets.UTF_8)) + "\",\"bytes\":" + body.length() + ",\"depth\":0,\"segments\":" + segments + "}}]}\n";
+ }
+
@Test
public void aReferenceTheSessionDoesNotHoldIsRefused() {
final SessionDataFile f = SessionDataFile.parse((
diff --git a/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/WithholdingShapesTest.java b/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/WithholdingShapesTest.java
index d074d4109c..b513e29481 100644
--- a/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/WithholdingShapesTest.java
+++ b/oap-server/analyzer/ai-agent-conversation/src/test/java/org/apache/skywalking/oap/server/ai/agent/conversation/withhold/WithholdingShapesTest.java
@@ -78,6 +78,14 @@ public class WithholdingShapesTest {
"{\"system\":[{\"type\":\"text\",\"content\":\"" + SECRET + "\"}],\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":[{\"name\":\"t\",\"description\":\"" + TOOL_SECRET + "\"}]}",
// Bedrock: the tools under the client's toolConfig, each in a toolSpec
"{\"system\":[{\"text\":\"" + SECRET + "\"}],\"messages\":[{\"role\":\"user\",\"content\":[{\"text\":\"" + PERSON + "\"}]}],\"toolConfig\":{\"tools\":[{\"toolSpec\":{\"name\":\"t\",\"description\":\"" + TOOL_SECRET + "\",\"inputSchema\":{\"json\":{}}}}]}}",
+ // a prompt key inside a set of tools keyed by name, beside a tool
+ "{\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":{\"t\":{\"description\":\"" + TOOL_SECRET + "\"},\"system_prompt\":\"" + SECRET + "\"}}",
+ // a set of tools that is, as a whole, a list of messages with LangChain's system pair in it
+ "{\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":[{\"name\":\"t\",\"description\":\"" + TOOL_SECRET + "\"},[\"system\",\"" + SECRET + "\"]]}",
+ // a set of tools keyed by name that is, as a whole, a system message
+ "{\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":{\"t\":{\"description\":\"" + TOOL_SECRET + "\"},\"role\":\"system\",\"content\":\"" + SECRET + "\"}}",
+ // a list of tools inside the set
+ "{\"system\":\"" + SECRET + "\",\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":[{\"name\":\"a\"},[{\"name\":\"t\",\"description\":\"" + TOOL_SECRET + "\"}]]}",
// LangChain JS: a serialized system message says its role only in the class that ends its id
"{\"messages\":[[{\"lc\":1,\"type\":\"constructor\",\"id\":[\"langchain_core\",\"messages\",\"SystemMessage\"],\"kwargs\":{\"content\":\"" + SECRET + "\",\"additional_kwargs\":{}}},{\"lc\":1,\"type\":\"constructor\",\"id\":[\"langchain_core\",\"messages\",\"HumanMessage\"],\"kwargs\":{\"content\":\"" + PERSON + "\"}}]],\"tools\":[{\"name\":\"t\",\"description\":\"" + TOOL_SECRET + "\"}]}",
})
@@ -519,30 +527,212 @@ public class WithholdingShapesTest {
public void aRequestKeptWholeInAnUnknownPartIsMasked(final boolean base64) {
final String body = "{\"system\":\"" + SECRET + "\",\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}]}";
final byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
+ final byte[] stored = bodyFile(keptWhole("r1", Manifest.ROLE_REQUEST, bytes, base64));
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ assertEquals(body, assertRebuilds(store, "r1"));
+ final String served = new String(Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes(), StandardCharsets.UTF_8);
+ assertFalse(served.contains(SECRET), served);
+ assertFalse(served.contains(Base64.getEncoder().encodeToString(bytes)), served);
+ final BodyStore again = new BodyStore();
+ again.add(SessionDataFile.parse(served.getBytes(StandardCharsets.UTF_8)));
+ final String masked = assertRebuilds(again, "r1");
+ assertTrue(masked.contains("[withheld: system_prompt]") && masked.contains(PERSON), masked);
+ }
+
+ /**
+ * A body that is not UTF-8 text cannot be written as a literal byte for byte, so it is withheld whole, never
+ * served with its stored digest over text that no longer has it: a request kept whole, its system prompt beside
+ * the byte FF, which is not one JSON object in UTF-8 text, and a response that copies all of it.
+ */
+ @Test
+ public void aBodyThatIsNotUtf8IsWithheldWholeRatherThanRewritten() {
+ final byte[] head = ("{\"system\":\"" + SECRET + "\",\"x\":\"").getBytes(StandardCharsets.UTF_8);
+ final byte[] bytes = Arrays.copyOf(head, head.length + 3);
+ bytes[head.length] = (byte) 0xff;
+ bytes[head.length + 1] = '"';
+ bytes[head.length + 2] = '}';
+ final String digest = Digests.sha256Hex(bytes);
+ final String response = "{\"ord\":2,\"off\":0,\"sha\":\"0\",\"bytes\":1,\"id\":\"b\",\"parts\":[{\"k\":\"data\",\"data\":"
+ + "{\"schema\":\"provider_body/1\",\"role\":\"response\",\"src\":\"b.json\",\"sha256\":\"" + digest
+ + "\",\"bytes\":" + bytes.length + ",\"depth\":1,\"segments\":[{\"copy\":{\"from\":\"a\",\"sha256\":\"" + digest
+ + "\",\"len\":" + bytes.length + "}}]},\"state\":\"available\",\"bytes\":1}]}";
+ final byte[] stored = bodyFile(keptWhole("a", Manifest.ROLE_REQUEST, bytes, true), response);
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ final byte[] served = Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes();
+ assertFalse(new String(served, StandardCharsets.UTF_8).contains(SECRET));
+ final SessionDataFile after = SessionDataFile.parse(served);
+ assertNull(Manifest.of(after.getRecords().get(0)), "a request that is not UTF-8 text is withheld whole");
+ final BodyStore again = new BodyStore();
+ again.add(after);
+ for (final SessionDataFile.Record rec : after.getRecords()) {
+ if (Manifest.of(rec) != null) {
+ // whatever is served with a manifest rebuilds to it
+ assertRebuilds(again, rec.getId());
+ continue;
+ }
+ for (final SessionDataFile.Part p : rec.getParts()) {
+ assertEquals("omitted", p.getState(), rec.getId());
+ }
+ }
+ }
+
+ /**
+ * A record landed again under an id the session holds, with another body, is not the one held, and a reader
+ * refuses it: it is withheld whole, so its own pieces, which the held body's masking knows nothing of, go too. A
+ * repeat with the same body is not the row held either, and is withheld whole as well; the row held is masked.
+ */
+ @Test
+ public void aRecordLandedAgainWithAnotherBodyIsWithheldWhole() {
+ final String held = "{\"system\":\"short\"}";
+ final String piece = "\"" + SECRET + "\"";
+ final String again = "{\"system\":" + piece + "}";
+ final byte[] stored = bodyFile(request("x", held, Collections.emptyList()),
+ record("x", Manifest.ROLE_REQUEST, again, Collections.singletonList(piece),
+ segments(lit("{\"system\":"), "{\"part\":0}", lit("}"))));
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ final String served = new String(Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes(), StandardCharsets.UTF_8);
+ assertFalse(served.contains(SECRET), served);
+ assertNull(Manifest.of(SessionDataFile.parse(served.getBytes(StandardCharsets.UTF_8)).getRecords().get(1)));
+
+ final byte[] twice = bodyFile(request("y", held, Collections.emptyList()), request("y", held, Collections.emptyList()));
+ final BodyStore store2 = new BodyStore();
+ store2.add(SessionDataFile.parse(twice));
+ final SessionDataFile after = SessionDataFile.parse(Withholding.file(twice, Collections.singletonList(Hide.SYSTEM_PROMPT), store2).getBytes());
+ assertNotNull(Manifest.of(after.getRecords().get(0)), "the row held is masked");
+ assertNull(Manifest.of(after.getRecords().get(1)), "a repeat with the same body is withheld whole");
+ }
+
+ /**
+ * A record the Sessionizer kept whole, its body in one unknown part with a why, as plain text or as base64.
+ */
+ private static String keptWhole(final String id, final String role, final byte[] bytes, final boolean base64) {
final JsonObject unknown = new JsonObject();
unknown.addProperty("k", "unknown");
unknown.addProperty("text", "it did not come back whole");
- unknown.addProperty("data", base64 ? Base64.getEncoder().encodeToString(bytes) : body);
+ unknown.addProperty("data", base64 ? Base64.getEncoder().encodeToString(bytes) : new String(bytes, StandardCharsets.UTF_8));
if (base64) {
unknown.addProperty("encoding", "base64");
}
unknown.addProperty("state", "available");
unknown.addProperty("bytes", bytes.length);
- final String manifest = "{\"schema\":\"provider_body/1\",\"role\":\"request\",\"src\":\"r1.json\",\"sha256\":\""
+ final String manifest = "{\"schema\":\"provider_body/1\",\"role\":\"" + role + "\",\"src\":\"" + id + ".json\",\"sha256\":\""
+ Digests.sha256Hex(bytes) + "\",\"bytes\":" + bytes.length + ",\"depth\":0,\"why\":\"it did not come back whole\",\"segments\":[{\"part\":0}]}";
- final String rec = "{\"ord\":1,\"off\":0,\"sha\":\"0\",\"bytes\":" + bytes.length + ",\"id\":\"r1\",\"parts\":[" + unknown
+ return "{\"ord\":1,\"off\":0,\"sha\":\"0\",\"bytes\":" + bytes.length + ",\"id\":\"" + id + "\",\"parts\":[" + unknown
+ ",{\"k\":\"data\",\"data\":" + manifest + ",\"state\":\"available\",\"bytes\":" + manifest.length() + "}]}";
- final byte[] stored = bodyFile(rec);
+ }
+
+ /**
+ * A response the session does not hold, here for a depth it claims with no copy, is withheld whole: it may refer
+ * to a piece of a request whose masked form no longer holds it, and a reader would not rebuild it anyway.
+ */
+ @Test
+ public void aResponseTheSessionDoesNotHoldIsWithheldWhole() {
+ final String piece = "\"" + SECRET + "\"";
+ final String requestBody = "{\"system\":" + piece + "}";
+ final String responseBody = "[" + piece + "]";
+ final String request = record("a", Manifest.ROLE_REQUEST, requestBody, Collections.singletonList(piece),
+ segments(lit("{\"system\":"), "{\"part\":0}", lit("}")));
+ final String response = record("b", Manifest.ROLE_RESPONSE, responseBody, Collections.emptyList(),
+ segments(lit("["), piece(piece), lit("]"))).replace("\"depth\":0", "\"depth\":1");
+ final byte[] stored = bodyFile(request, response);
final BodyStore store = new BodyStore();
store.add(SessionDataFile.parse(stored));
- assertEquals(body, assertRebuilds(store, "r1"));
- final String served = new String(Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes(), StandardCharsets.UTF_8);
- assertFalse(served.contains(SECRET), served);
- assertFalse(served.contains(Base64.getEncoder().encodeToString(bytes)), served);
+ final byte[] served = Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes();
+ assertFalse(new String(served, StandardCharsets.UTF_8).contains(SECRET));
+ final SessionDataFile after = SessionDataFile.parse(served);
+ assertNull(Manifest.of(after.getRecords().get(1)), "a response the session does not hold is withheld whole");
final BodyStore again = new BodyStore();
- again.add(SessionDataFile.parse(served.getBytes(StandardCharsets.UTF_8)));
+ again.add(after);
+ assertTrue(assertRebuilds(again, "a").contains("[withheld: system_prompt]"));
+ }
+
+ /**
+ * A row is served as the one its session holds, decided when it lands, so a later file cannot change it: a request
+ * in seq 1 with a depth it may not claim, and the same request again in seq 2, valid. Seq 1 is served the same
+ * whether or not seq 2 is read, its request withheld whole either way.
+ */
+ @Test
+ public void aLaterRecordUnderTheSameIdDoesNotChangeAnEarlierRow() {
+ final String body = "{\"system\":\"" + SECRET + "\"}";
+ final byte[] first = bodyFile(request("x", body, Collections.emptyList()).replace("\"depth\":0", "\"depth\":1"));
+ final byte[] second = new String(bodyFile(request("x", body, Collections.emptyList())), StandardCharsets.UTF_8)
+ .replace("\"seq\":1", "\"seq\":2").getBytes(StandardCharsets.UTF_8);
+ final BodyStore alone = new BodyStore();
+ alone.add(SessionDataFile.parse(first));
+ final BodyStore both = new BodyStore();
+ both.add(SessionDataFile.parse(first));
+ both.add(SessionDataFile.parse(second));
+ final byte[] servedAlone = Withholding.file(first, Collections.singletonList(Hide.SYSTEM_PROMPT), alone).getBytes();
+ final byte[] servedWithLater = Withholding.file(first, Collections.singletonList(Hide.SYSTEM_PROMPT), both).getBytes();
+ assertEquals(new String(servedAlone, StandardCharsets.UTF_8), new String(servedWithLater, StandardCharsets.UTF_8));
+ assertNull(Manifest.of(SessionDataFile.parse(servedAlone).getRecords().get(0)));
+ // the valid one in seq 2 is the one held, and is masked
+ final String laterServed = new String(Withholding.file(second, Collections.singletonList(Hide.SYSTEM_PROMPT), both).getBytes(), StandardCharsets.UTF_8);
+ assertFalse(laterServed.contains(SECRET), laterServed);
+ assertNotNull(Manifest.of(SessionDataFile.parse(laterServed.getBytes(StandardCharsets.UTF_8)).getRecords().get(0)));
+ }
+
+ /** A tool named like a prompt key keeps the tools' marker: a masked tool is not walked again. */
+ @Test
+ public void aToolNamedLikeAPromptKeyKeepsTheToolsMarker() {
+ final String body = "{\"messages\":[{\"role\":\"user\",\"content\":\"" + PERSON + "\"}],\"tools\":{\"system\":{\"description\":\"" + TOOL_SECRET + "\"}}}";
+ final byte[] stored = bodyFile(request("r1", body, Collections.emptyList()));
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ final byte[] served = Withholding.file(stored, BOTH, store).getBytes();
+ final BodyStore again = new BodyStore();
+ again.add(SessionDataFile.parse(served));
final String masked = assertRebuilds(again, "r1");
- assertTrue(masked.contains("[withheld: system_prompt]") && masked.contains(PERSON), masked);
+ assertTrue(masked.contains("\"system\":{\"withheld\":\"tool_schemas\"}"), masked);
+ assertFalse(masked.contains(MARKER_SYSTEM_TEXT), masked);
+ }
+
+ private static final String MARKER_SYSTEM_TEXT = "[withheld: system_prompt]";
+
+ /**
+ * A body of another role holds no piece a later body may use: it is not held, so a response that uses its piece is
+ * not held either, and both are withheld whole, rather than the response served referring to a piece that is gone.
+ */
+ @ParameterizedTest
+ @ValueSource(strings = {"\"REQUEST\"", "5"})
+ public void aPieceOfABodyOfAnotherRoleIsUsedByNothingServed(final String role) {
+ final String piece = "\"x\"";
+ final String holder = record("a", Manifest.ROLE_REQUEST, piece, Collections.singletonList(piece), segments("{\"part\":0}"))
+ .replace("\"role\":\"request\"", "\"role\":" + role);
+ final String user = record("b", Manifest.ROLE_RESPONSE, piece, Collections.emptyList(), segments(piece(piece)));
+ final byte[] stored = bodyFile(holder, user);
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ final SessionDataFile after = SessionDataFile.parse(Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes());
+ final BodyStore again = new BodyStore();
+ again.add(after);
+ for (final SessionDataFile.Record rec : after.getRecords()) {
+ if (Manifest.of(rec) != null) {
+ // whatever is served with a manifest rebuilds from the served file alone
+ assertRebuilds(again, rec.getId());
+ }
+ }
+ assertNull(Manifest.of(after.getRecords().get(0)));
+ assertNull(Manifest.of(after.getRecords().get(1)));
+ }
+
+ /**
+ * A body whose role is neither request nor response, as the page gives it, may be a request, so it is withheld
+ * whole rather than served as stored like a response.
+ */
+ @ParameterizedTest
+ @ValueSource(strings = {"\"REQUEST\"", "5", "\"\""})
+ public void aBodyOfAnotherRoleIsWithheldWhole(final String role) {
+ final String body = "{\"system\":\"" + SECRET + "\"}";
+ final byte[] stored = bodyFile(request("r1", body, Collections.emptyList()).replace("\"role\":\"request\"", "\"role\":" + role));
+ final BodyStore store = new BodyStore();
+ store.add(SessionDataFile.parse(stored));
+ final String served = new String(Withholding.file(stored, Collections.singletonList(Hide.SYSTEM_PROMPT), store).getBytes(), StandardCharsets.UTF_8);
+ assertFalse(served.contains(SECRET), served);
+ assertNull(Manifest.of(SessionDataFile.parse(served.getBytes(StandardCharsets.UTF_8)).getRecords().get(0)));
}
private static String piece(final String bytes) {