Skip to content

Commit 067496b

Browse files
vincentkocvincentkoc
committed
refactor(providers): share anthropic payload policy
1 parent 3e0ddaf commit 067496b

7 files changed

Lines changed: 536 additions & 118 deletions

File tree

‎extensions/anthropic/stream-wrappers.ts‎

Lines changed: 25 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import type { StreamFn } from "@mariozechner/pi-agent-core";
22
import { streamSimple } from "@mariozechner/pi-ai";
3-
import { resolveProviderRequestCapabilities } from "openclaw/plugin-sdk/provider-http";
4-
import { streamWithPayloadPatch } from "openclaw/plugin-sdk/provider-stream";
3+
import {
4+
applyAnthropicPayloadPolicyToParams,
5+
resolveAnthropicPayloadPolicy,
6+
streamWithPayloadPatch,
7+
} from "openclaw/plugin-sdk/provider-stream";
58
import { createSubsystemLogger } from "openclaw/plugin-sdk/runtime-env";
69

710
const log = createSubsystemLogger("anthropic-stream");
@@ -52,20 +55,6 @@ function isAnthropicOAuthApiKey(apiKey: unknown): boolean {
5255
return typeof apiKey === "string" && apiKey.includes("sk-ant-oat");
5356
}
5457

55-
function allowsAnthropicServiceTier(model: {
56-
api?: unknown;
57-
provider?: unknown;
58-
baseUrl?: unknown;
59-
}): boolean {
60-
return resolveProviderRequestCapabilities({
61-
provider: typeof model.provider === "string" ? model.provider : undefined,
62-
api: typeof model.api === "string" ? model.api : undefined,
63-
baseUrl: typeof model.baseUrl === "string" ? model.baseUrl : undefined,
64-
capability: "llm",
65-
transport: "stream",
66-
}).allowsAnthropicServiceTier;
67-
}
68-
6958
function resolveAnthropicFastServiceTier(enabled: boolean): AnthropicServiceTier {
7059
return enabled ? "auto" : "standard_only";
7160
}
@@ -161,15 +150,19 @@ export function createAnthropicFastModeWrapper(
161150
const underlying = baseStreamFn ?? streamSimple;
162151
const serviceTier = resolveAnthropicFastServiceTier(enabled);
163152
return (model, context, options) => {
164-
if (!allowsAnthropicServiceTier(model)) {
153+
const payloadPolicy = resolveAnthropicPayloadPolicy({
154+
provider: typeof model.provider === "string" ? model.provider : undefined,
155+
api: typeof model.api === "string" ? model.api : undefined,
156+
baseUrl: typeof model.baseUrl === "string" ? model.baseUrl : undefined,
157+
serviceTier,
158+
});
159+
if (!payloadPolicy.allowsServiceTier) {
165160
return underlying(model, context, options);
166161
}
167162

168-
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) => {
169-
if (payloadObj.service_tier === undefined) {
170-
payloadObj.service_tier = serviceTier;
171-
}
172-
});
163+
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) =>
164+
applyAnthropicPayloadPolicyToParams(payloadObj, payloadPolicy),
165+
);
173166
};
174167
}
175168

@@ -179,15 +172,19 @@ export function createAnthropicServiceTierWrapper(
179172
): StreamFn {
180173
const underlying = baseStreamFn ?? streamSimple;
181174
return (model, context, options) => {
182-
if (!allowsAnthropicServiceTier(model)) {
175+
const payloadPolicy = resolveAnthropicPayloadPolicy({
176+
provider: typeof model.provider === "string" ? model.provider : undefined,
177+
api: typeof model.api === "string" ? model.api : undefined,
178+
baseUrl: typeof model.baseUrl === "string" ? model.baseUrl : undefined,
179+
serviceTier,
180+
});
181+
if (!payloadPolicy.allowsServiceTier) {
183182
return underlying(model, context, options);
184183
}
185184

186-
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) => {
187-
if (payloadObj.service_tier === undefined) {
188-
payloadObj.service_tier = serviceTier;
189-
}
190-
});
185+
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) =>
186+
applyAnthropicPayloadPolicyToParams(payloadObj, payloadPolicy),
187+
);
191188
};
192189
}
193190

‎src/agents/anthropic-payload-policy.test.ts‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import { describe, expect, it } from "vitest";
2+
import {
3+
applyAnthropicPayloadPolicyToParams,
4+
resolveAnthropicPayloadPolicy,
5+
} from "./anthropic-payload-policy.js";
6+
7+
describe("anthropic payload policy", () => {
8+
it("applies native Anthropic service tier and cache markers without widening cache scope", () => {
9+
const policy = resolveAnthropicPayloadPolicy({
10+
provider: "anthropic",
11+
api: "anthropic-messages",
12+
baseUrl: "https://api.anthropic.com/v1",
13+
cacheRetention: "long",
14+
enableCacheControl: true,
15+
serviceTier: "standard_only",
16+
});
17+
const payload: Record<string, unknown> = {
18+
system: [
19+
{ type: "text", text: "Follow policy." },
20+
{ type: "text", text: "Use tools carefully." },
21+
],
22+
messages: [
23+
{
24+
role: "assistant",
25+
content: [{ type: "text", text: "Working." }],
26+
},
27+
{
28+
role: "user",
29+
content: [
30+
{ type: "text", text: "Hello" },
31+
{ type: "tool_result", tool_use_id: "tool_1", content: "done" },
32+
],
33+
},
34+
],
35+
};
36+
37+
applyAnthropicPayloadPolicyToParams(payload, policy);
38+
39+
expect(payload.service_tier).toBe("standard_only");
40+
expect(payload.system).toEqual([
41+
{
42+
type: "text",
43+
text: "Follow policy.",
44+
cache_control: { type: "ephemeral", ttl: "1h" },
45+
},
46+
{
47+
type: "text",
48+
text: "Use tools carefully.",
49+
cache_control: { type: "ephemeral", ttl: "1h" },
50+
},
51+
]);
52+
expect(payload.messages[0]).toEqual({
53+
role: "assistant",
54+
content: [{ type: "text", text: "Working." }],
55+
});
56+
expect(payload.messages[1]).toEqual({
57+
role: "user",
58+
content: [
59+
{ type: "text", text: "Hello" },
60+
{
61+
type: "tool_result",
62+
tool_use_id: "tool_1",
63+
content: "done",
64+
cache_control: { type: "ephemeral", ttl: "1h" },
65+
},
66+
],
67+
});
68+
});
69+
70+
it("denies proxied Anthropic service tier and omits long-TTL upgrades for custom hosts", () => {
71+
const policy = resolveAnthropicPayloadPolicy({
72+
provider: "anthropic",
73+
api: "anthropic-messages",
74+
baseUrl: "https://proxy.example.com/anthropic",
75+
cacheRetention: "long",
76+
enableCacheControl: true,
77+
serviceTier: "auto",
78+
});
79+
const payload: Record<string, unknown> = {
80+
system: [{ type: "text", text: "Follow policy." }],
81+
messages: [{ role: "user", content: "Hello" }],
82+
};
83+
84+
applyAnthropicPayloadPolicyToParams(payload, policy);
85+
86+
expect(payload).not.toHaveProperty("service_tier");
87+
expect(payload.system).toEqual([
88+
{
89+
type: "text",
90+
text: "Follow policy.",
91+
cache_control: { type: "ephemeral" },
92+
},
93+
]);
94+
expect(payload.messages[0]).toEqual({
95+
role: "user",
96+
content: [{ type: "text", text: "Hello", cache_control: { type: "ephemeral" } }],
97+
});
98+
});
99+
});

‎src/agents/anthropic-payload-policy.ts‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
import { resolveProviderRequestCapabilities } from "./provider-attribution.js";
2+
3+
export type AnthropicServiceTier = "auto" | "standard_only";
4+
5+
export type AnthropicEphemeralCacheControl = {
6+
type: "ephemeral";
7+
ttl?: "1h";
8+
};
9+
10+
type AnthropicPayloadPolicyInput = {
11+
api?: string;
12+
baseUrl?: string;
13+
cacheRetention?: "short" | "long" | "none";
14+
enableCacheControl?: boolean;
15+
provider?: string;
16+
serviceTier?: AnthropicServiceTier;
17+
};
18+
19+
export type AnthropicPayloadPolicy = {
20+
allowsServiceTier: boolean;
21+
cacheControl: AnthropicEphemeralCacheControl | undefined;
22+
serviceTier: AnthropicServiceTier | undefined;
23+
};
24+
25+
function resolveAnthropicEphemeralCacheControl(
26+
baseUrl: string | undefined,
27+
cacheRetention: AnthropicPayloadPolicyInput["cacheRetention"],
28+
): AnthropicEphemeralCacheControl | undefined {
29+
const retention =
30+
cacheRetention ?? (process.env.PI_CACHE_RETENTION === "long" ? "long" : "short");
31+
if (retention === "none") {
32+
return undefined;
33+
}
34+
const ttl =
35+
retention === "long" && typeof baseUrl === "string" && baseUrl.includes("api.anthropic.com")
36+
? "1h"
37+
: undefined;
38+
return { type: "ephemeral", ...(ttl ? { ttl } : {}) };
39+
}
40+
41+
function applyAnthropicCacheControlToSystem(
42+
system: unknown,
43+
cacheControl: AnthropicEphemeralCacheControl,
44+
): void {
45+
if (!Array.isArray(system)) {
46+
return;
47+
}
48+
49+
for (const block of system) {
50+
if (!block || typeof block !== "object") {
51+
continue;
52+
}
53+
const record = block as Record<string, unknown>;
54+
if (record.type === "text" && record.cache_control === undefined) {
55+
record.cache_control = cacheControl;
56+
}
57+
}
58+
}
59+
60+
function applyAnthropicCacheControlToMessages(
61+
messages: unknown,
62+
cacheControl: AnthropicEphemeralCacheControl,
63+
): void {
64+
if (!Array.isArray(messages) || messages.length === 0) {
65+
return;
66+
}
67+
68+
const lastMessage = messages[messages.length - 1];
69+
if (!lastMessage || typeof lastMessage !== "object") {
70+
return;
71+
}
72+
73+
const record = lastMessage as Record<string, unknown>;
74+
if (record.role !== "user") {
75+
return;
76+
}
77+
78+
const content = record.content;
79+
if (Array.isArray(content)) {
80+
const lastBlock = content[content.length - 1];
81+
if (!lastBlock || typeof lastBlock !== "object") {
82+
return;
83+
}
84+
const lastBlockRecord = lastBlock as Record<string, unknown>;
85+
if (
86+
lastBlockRecord.type === "text" ||
87+
lastBlockRecord.type === "image" ||
88+
lastBlockRecord.type === "tool_result"
89+
) {
90+
lastBlockRecord.cache_control = cacheControl;
91+
}
92+
return;
93+
}
94+
95+
if (typeof content === "string") {
96+
record.content = [
97+
{
98+
type: "text",
99+
text: content,
100+
cache_control: cacheControl,
101+
},
102+
];
103+
}
104+
}
105+
106+
export function resolveAnthropicPayloadPolicy(
107+
input: AnthropicPayloadPolicyInput,
108+
): AnthropicPayloadPolicy {
109+
const capabilities = resolveProviderRequestCapabilities({
110+
provider: input.provider,
111+
api: input.api,
112+
baseUrl: input.baseUrl,
113+
capability: "llm",
114+
transport: "stream",
115+
});
116+
117+
return {
118+
allowsServiceTier: capabilities.allowsAnthropicServiceTier,
119+
cacheControl:
120+
input.enableCacheControl === true
121+
? resolveAnthropicEphemeralCacheControl(input.baseUrl, input.cacheRetention)
122+
: undefined,
123+
serviceTier: input.serviceTier,
124+
};
125+
}
126+
127+
export function applyAnthropicPayloadPolicyToParams(
128+
payloadObj: Record<string, unknown>,
129+
policy: AnthropicPayloadPolicy,
130+
): void {
131+
if (
132+
policy.allowsServiceTier &&
133+
policy.serviceTier !== undefined &&
134+
payloadObj.service_tier === undefined
135+
) {
136+
payloadObj.service_tier = policy.serviceTier;
137+
}
138+
139+
if (!policy.cacheControl) {
140+
return;
141+
}
142+
143+
applyAnthropicCacheControlToSystem(payloadObj.system, policy.cacheControl);
144+
// Preserve Anthropic cache-write scope by only tagging the trailing user turn.
145+
applyAnthropicCacheControlToMessages(payloadObj.messages, policy.cacheControl);
146+
}
147+
148+
export function applyAnthropicEphemeralCacheControlMarkers(
149+
payloadObj: Record<string, unknown>,
150+
): void {
151+
const messages = payloadObj.messages;
152+
if (!Array.isArray(messages)) {
153+
return;
154+
}
155+
156+
for (const message of messages as Array<{ role?: string; content?: unknown }>) {
157+
if (message.role === "system" || message.role === "developer") {
158+
if (typeof message.content === "string") {
159+
message.content = [
160+
{ type: "text", text: message.content, cache_control: { type: "ephemeral" } },
161+
];
162+
continue;
163+
}
164+
if (Array.isArray(message.content) && message.content.length > 0) {
165+
const last = message.content[message.content.length - 1];
166+
if (last && typeof last === "object") {
167+
const record = last as Record<string, unknown>;
168+
if (record.type !== "thinking" && record.type !== "redacted_thinking") {
169+
record.cache_control = { type: "ephemeral" };
170+
}
171+
}
172+
}
173+
continue;
174+
}
175+
176+
if (message.role === "assistant" && Array.isArray(message.content)) {
177+
for (const block of message.content) {
178+
if (!block || typeof block !== "object") {
179+
continue;
180+
}
181+
const record = block as Record<string, unknown>;
182+
if (record.type === "thinking" || record.type === "redacted_thinking") {
183+
delete record.cache_control;
184+
}
185+
}
186+
}
187+
}
188+
}

0 commit comments

Comments
 (0)