Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 27fea03c07 | |||
| 250af4df36 | |||
| 884bb8c09f | |||
| 0c152a24d2 | |||
| 3345544921 | |||
| 8e2597c877 | |||
| d241dd7f9e | |||
| d437c31da4 |
@@ -36,6 +36,20 @@ interface A2AResponseShape {
|
||||
error?: { message?: string };
|
||||
}
|
||||
|
||||
// Wire shape for GET /workspaces/:id/chat-history (chat_history.go → ChatHistoryResponse).
|
||||
interface ApiChatMessage {
|
||||
id: string;
|
||||
role: string; // "user" | "agent" | "system"
|
||||
content: string;
|
||||
timestamp: string;
|
||||
attachments?: Array<{ name: string; uri: string; mimeType?: string; size?: number }>;
|
||||
}
|
||||
|
||||
interface ChatHistoryResponse {
|
||||
messages: ApiChatMessage[];
|
||||
reached_end: boolean;
|
||||
}
|
||||
|
||||
const formatTime = (date: Date) =>
|
||||
date.toLocaleTimeString([], { hour: "numeric", minute: "2-digit" });
|
||||
|
||||
@@ -61,18 +75,14 @@ export function MobileChat({
|
||||
// that creates a new [] reference on every store update when the key is
|
||||
// absent, causing infinite re-render (React error #185).
|
||||
const storedMessages = useCanvasStore((s) => s.agentMessages[agentId]);
|
||||
const [messages, setMessages] = useState<ChatMessage[]>(() =>
|
||||
(storedMessages ?? []).map((m) => ({
|
||||
id: m.id,
|
||||
role: "agent",
|
||||
text: m.content,
|
||||
ts: formatStoredTimestamp(m.timestamp),
|
||||
})),
|
||||
);
|
||||
// Start empty — history is loaded via useEffect below.
|
||||
const [messages, setMessages] = useState<ChatMessage[]>([]);
|
||||
const [draft, setDraft] = useState("");
|
||||
const [tab, setTab] = useState<SubTab>("my");
|
||||
const [sending, setSending] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [loading, setLoading] = useState(true); // history is loading on mount
|
||||
const [historyError, setHistoryError] = useState<string | null>(null);
|
||||
const scrollRef = useRef<HTMLDivElement>(null);
|
||||
// Synchronous re-entry guard. `setSending(true)` schedules a state
|
||||
// update but doesn't flush before a second tap can fire send() — a ref
|
||||
@@ -80,6 +90,9 @@ export function MobileChat({
|
||||
// double-send race a stale `sending` lets through.
|
||||
const sendInFlightRef = useRef(false);
|
||||
const composerRef = useRef<HTMLTextAreaElement>(null);
|
||||
// Guard: don't treat the initial store population as a live push.
|
||||
// Set to false after the first render completes.
|
||||
const initDoneRef = useRef(false);
|
||||
|
||||
// Auto-grow the textarea: reset height to 'auto' so the scrollHeight
|
||||
// shrinks when the user deletes text, then size to scrollHeight up to
|
||||
@@ -92,6 +105,75 @@ export function MobileChat({
|
||||
el.style.height = `${next}px`;
|
||||
}, [draft]);
|
||||
|
||||
// Fetch chat history on mount; keep merging live agentMessages while the
|
||||
// panel is open. InitDoneRef prevents the initial store snapshot from
|
||||
// triggering the live-merge path (the store buffer is populated by
|
||||
// ChatTab on desktop, not on mobile — this effect loads history as the
|
||||
// mobile-native path).
|
||||
useEffect(() => {
|
||||
let cancelled = false;
|
||||
|
||||
const mapApiMessage = (m: ApiChatMessage): ChatMessage => ({
|
||||
id: m.id,
|
||||
role: m.role === "user" ? "user" : "agent",
|
||||
text: m.content,
|
||||
ts: formatStoredTimestamp(m.timestamp),
|
||||
});
|
||||
|
||||
const syncLive = () => {
|
||||
const live = useCanvasStore.getState().agentMessages[agentId] ?? [];
|
||||
if (live.length > 0) {
|
||||
setMessages((prev) => {
|
||||
const existingIds = new Set(prev.map((m) => m.id));
|
||||
const newOnes = live
|
||||
.filter((m) => !existingIds.has(m.id))
|
||||
.map((m) => ({
|
||||
id: m.id,
|
||||
role: "agent" as const,
|
||||
text: m.content,
|
||||
ts: formatStoredTimestamp(m.timestamp),
|
||||
}));
|
||||
return newOnes.length > 0 ? [...prev, ...newOnes] : prev;
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
const bootstrap = async (): Promise<(() => void) | undefined> => {
|
||||
setLoading(true);
|
||||
setHistoryError(null);
|
||||
try {
|
||||
const res = await api.get<ChatHistoryResponse>(
|
||||
`/workspaces/${agentId}/chat-history?limit=50`,
|
||||
);
|
||||
if (cancelled) return;
|
||||
const initial = (res.messages ?? []).map(mapApiMessage);
|
||||
setMessages(initial);
|
||||
// Mark init done BEFORE marking loading=false so any store push
|
||||
// that arrives in the same tick is treated as live, not init.
|
||||
initDoneRef.current = true;
|
||||
setLoading(false);
|
||||
// Subscribe to live pushes after init is complete.
|
||||
syncLive();
|
||||
const unsubscribe = useCanvasStore.subscribe(syncLive);
|
||||
return unsubscribe; // returned for cleanup
|
||||
} catch (e) {
|
||||
if (cancelled) return;
|
||||
setHistoryError(e instanceof Error ? e.message : "Failed to load chat history");
|
||||
setLoading(false);
|
||||
initDoneRef.current = true;
|
||||
return undefined;
|
||||
}
|
||||
};
|
||||
|
||||
let maybeUnsubscribe: (() => void) | undefined;
|
||||
bootstrap().then((fn) => { maybeUnsubscribe = fn; });
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
if (maybeUnsubscribe) maybeUnsubscribe();
|
||||
};
|
||||
}, [agentId]);
|
||||
|
||||
useEffect(() => {
|
||||
if (scrollRef.current) {
|
||||
scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
|
||||
@@ -311,7 +393,61 @@ export function MobileChat({
|
||||
Agent Comms — peer-to-peer A2A traffic surfaces in the Comms tab.
|
||||
</div>
|
||||
)}
|
||||
{tab === "my" && messages.length === 0 && (
|
||||
{tab === "my" && loading && (
|
||||
<div style={{ padding: "20px 4px", textAlign: "center", color: p.text3, fontSize: 13 }}>
|
||||
<div style={{ marginBottom: 6, opacity: 0.6, animation: "spin 1s linear infinite", display: "inline-block", fontSize: 16 }}>⟳</div>
|
||||
<div>Loading chat history…</div>
|
||||
</div>
|
||||
)}
|
||||
{tab === "my" && !loading && historyError && (
|
||||
<div
|
||||
role="alert"
|
||||
style={{
|
||||
padding: "14px 4px",
|
||||
textAlign: "center",
|
||||
color: p.failed,
|
||||
fontSize: 13,
|
||||
}}
|
||||
>
|
||||
<div style={{ marginBottom: 8 }}>Could not load chat history.</div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setLoading(true);
|
||||
setHistoryError(null);
|
||||
api.get(`/workspaces/${agentId}/chat-history?limit=50`).then(
|
||||
(res: unknown) => {
|
||||
const r = res as ChatHistoryResponse;
|
||||
setMessages((r.messages ?? []).map((m) => ({
|
||||
id: m.id,
|
||||
role: m.role === "user" ? "user" : "agent",
|
||||
text: m.content,
|
||||
ts: formatStoredTimestamp(m.timestamp),
|
||||
})));
|
||||
setLoading(false);
|
||||
initDoneRef.current = true;
|
||||
},
|
||||
).catch((e: unknown) => {
|
||||
setHistoryError(e instanceof Error ? e.message : "Failed to load");
|
||||
setLoading(false);
|
||||
initDoneRef.current = true;
|
||||
});
|
||||
}}
|
||||
style={{
|
||||
padding: "6px 14px",
|
||||
borderRadius: 14,
|
||||
border: `0.5px solid ${p.failed}`,
|
||||
background: "transparent",
|
||||
color: p.failed,
|
||||
fontSize: 12,
|
||||
cursor: "pointer",
|
||||
}}
|
||||
>
|
||||
Retry
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
{tab === "my" && !loading && !historyError && messages.length === 0 && (
|
||||
<div style={{ padding: "20px 4px", textAlign: "center", color: p.text3, fontSize: 13 }}>
|
||||
Send a message to start chatting.
|
||||
</div>
|
||||
|
||||
@@ -8,11 +8,19 @@
|
||||
* NOTE: No @testing-library/jest-dom — use DOM APIs.
|
||||
*/
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { cleanup, render } from "@testing-library/react";
|
||||
import { act, cleanup, render, waitFor } from "@testing-library/react";
|
||||
import React from "react";
|
||||
|
||||
import { MobileChat } from "../MobileChat";
|
||||
|
||||
// ─── Mock API ─────────────────────────────────────────────────────────────────
|
||||
// vi.mock without a factory auto-mocks the module. In tests, we configure
|
||||
// api.get / api.post directly (they are vi.fn() from the auto-mock).
|
||||
// Tests that need specific behaviour use mockResolvedValueOnce on the
|
||||
// auto-mocked functions.
|
||||
vi.mock("@/lib/api");
|
||||
import { api } from "@/lib/api";
|
||||
|
||||
// ─── Mock store ───────────────────────────────────────────────────────────────
|
||||
|
||||
const mockAgentId = "ws-chat-test";
|
||||
@@ -32,8 +40,14 @@ const mockStoreState = {
|
||||
|
||||
vi.mock("@/store/canvas", () => ({
|
||||
useCanvasStore: Object.assign(
|
||||
vi.fn((sel) => sel(mockStoreState)),
|
||||
{ getState: () => mockStoreState },
|
||||
vi.fn((sel?: (state: typeof mockStoreState) => unknown) => {
|
||||
if (sel) return sel(mockStoreState);
|
||||
return mockStoreState;
|
||||
}),
|
||||
{
|
||||
getState: () => mockStoreState,
|
||||
subscribe: vi.fn(() => vi.fn()),
|
||||
},
|
||||
),
|
||||
summarizeWorkspaceCapabilities: vi.fn((data: Record<string, unknown>) => {
|
||||
const agentCard = data.agentCard as Record<string, unknown> | null;
|
||||
@@ -54,16 +68,6 @@ vi.mock("@/store/canvas", () => ({
|
||||
}),
|
||||
}));
|
||||
|
||||
// ─── Mock API ─────────────────────────────────────────────────────────────────
|
||||
|
||||
const { mockApiPost } = vi.hoisted(() => ({
|
||||
mockApiPost: vi.fn().mockResolvedValue({ result: { parts: [] } }),
|
||||
}));
|
||||
|
||||
vi.mock("@/lib/api", () => ({
|
||||
api: { post: mockApiPost },
|
||||
}));
|
||||
|
||||
// ─── Fixtures ────────────────────────────────────────────────────────────────
|
||||
|
||||
const onlineNode = {
|
||||
@@ -150,7 +154,15 @@ beforeEach(() => {
|
||||
mockOnBack.mockClear();
|
||||
mockStoreState.nodes = [];
|
||||
mockStoreState.agentMessages = {};
|
||||
mockApiPost.mockClear();
|
||||
// Set up spies on the real api methods. Tests override these per-call.
|
||||
const getSpy = vi.spyOn(api, "get");
|
||||
const postSpy = vi.spyOn(api, "post");
|
||||
getSpy.mockResolvedValue({ messages: [], reached_end: true });
|
||||
postSpy.mockResolvedValue({ result: { parts: [] } });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
@@ -266,15 +278,26 @@ describe("MobileChat — empty state", () => {
|
||||
mockStoreState.nodes = [onlineNode];
|
||||
});
|
||||
|
||||
it('shows "Send a message to start chatting." when no messages', () => {
|
||||
const { container } = renderChat(mockAgentId);
|
||||
it('shows "Send a message to start chatting." when no messages', async () => {
|
||||
// History fetch resolves immediately in tests (mockResolvedValue).
|
||||
// act() flushes the microtask queue so the component reaches its
|
||||
// post-load state before we assert.
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("Send a message to start chatting.");
|
||||
});
|
||||
|
||||
it("shows no messages when agentMessages[agentId] is absent (undefined)", () => {
|
||||
it("shows no messages when agentMessages[agentId] is absent (undefined)", async () => {
|
||||
// Explicitly set to empty to simulate no stored messages
|
||||
mockStoreState.agentMessages = {};
|
||||
const { container } = renderChat(mockAgentId);
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("Send a message to start chatting.");
|
||||
});
|
||||
});
|
||||
@@ -321,3 +344,132 @@ describe("MobileChat — dark mode", () => {
|
||||
expect(container.querySelector('[aria-label="Back"]')).toBeTruthy();
|
||||
});
|
||||
});
|
||||
|
||||
// ─── Chat history loading ────────────────────────────────────────────────────
|
||||
|
||||
describe("MobileChat — chat history", () => {
|
||||
beforeEach(() => {
|
||||
mockStoreState.nodes = [onlineNode];
|
||||
});
|
||||
|
||||
it("calls GET /workspaces/:id/chat-history on mount", async () => {
|
||||
await act(async () => {
|
||||
renderChat(mockAgentId);
|
||||
});
|
||||
expect(api.get).toHaveBeenCalledWith(
|
||||
`/workspaces/${mockAgentId}/chat-history?limit=50`,
|
||||
);
|
||||
});
|
||||
|
||||
it("shows loading state while history is fetching", () => {
|
||||
// Do NOT await — check the pre-resolve state.
|
||||
const { container } = renderChat(mockAgentId);
|
||||
expect(container.textContent ?? "").toContain("Loading chat history…");
|
||||
});
|
||||
|
||||
it("shows empty state after history resolves with no messages", async () => {
|
||||
// beforeEach already sets api.get to resolve with empty — no override needed.
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("Send a message to start chatting.");
|
||||
});
|
||||
|
||||
it("renders messages from history response", async () => {
|
||||
vi.spyOn(api, "get").mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: "msg-1",
|
||||
role: "user",
|
||||
content: "Hello agent",
|
||||
timestamp: "2026-04-25T10:00:00Z",
|
||||
},
|
||||
{
|
||||
id: "msg-2",
|
||||
role: "agent",
|
||||
content: "Hello back",
|
||||
timestamp: "2026-04-25T10:00:01Z",
|
||||
},
|
||||
],
|
||||
reached_end: true,
|
||||
});
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("Hello agent");
|
||||
expect(container.textContent ?? "").toContain("Hello back");
|
||||
});
|
||||
|
||||
it("maps user role from API correctly", async () => {
|
||||
vi.spyOn(api, "get").mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: "msg-u",
|
||||
role: "user",
|
||||
content: "user message",
|
||||
timestamp: "2026-04-25T10:00:00Z",
|
||||
},
|
||||
],
|
||||
reached_end: true,
|
||||
});
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
// User messages render right-aligned. The text content check is sufficient
|
||||
// to confirm the message appeared.
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("user message");
|
||||
});
|
||||
|
||||
it("shows error state when history fetch fails", async () => {
|
||||
vi.spyOn(api, "get").mockRejectedValue(new Error("Network error"));
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
expect(container.textContent ?? "").toContain("Could not load chat history.");
|
||||
expect(container.textContent ?? "").toContain("Retry");
|
||||
});
|
||||
|
||||
it("Retry button re-fetches history after error", async () => {
|
||||
// Make the initial mount call fail so the Retry button appears, then
|
||||
// make the retry call succeed so we can verify the full flow.
|
||||
const getSpy = vi.spyOn(api, "get");
|
||||
getSpy
|
||||
.mockRejectedValueOnce(new Error("Network error"))
|
||||
.mockResolvedValueOnce({ messages: [], reached_end: true });
|
||||
|
||||
let renderResult: ReturnType<typeof renderChat>;
|
||||
await act(async () => {
|
||||
renderResult = renderChat(mockAgentId);
|
||||
});
|
||||
const { container } = renderResult!;
|
||||
|
||||
// Error state should be shown with Retry button.
|
||||
expect(container.textContent ?? "").toContain("Could not load chat history.");
|
||||
expect(container.textContent ?? "").toContain("Retry");
|
||||
|
||||
// Click Retry — the button's onClick fires api.get again.
|
||||
// The second mockResolvedValueOnce makes it succeed.
|
||||
const retryBtn = Array.from(container.querySelectorAll("button")).find(
|
||||
(b) => b.textContent?.trim() === "Retry",
|
||||
);
|
||||
expect(retryBtn).toBeTruthy();
|
||||
await act(async () => {
|
||||
retryBtn?.click();
|
||||
});
|
||||
|
||||
// waitFor polls until the retry resolves and component re-renders.
|
||||
await waitFor(() => {
|
||||
expect(container.textContent ?? "").toContain("Send a message to start chatting.");
|
||||
});
|
||||
// Initial call + retry = 2.
|
||||
expect(getSpy).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -80,26 +80,103 @@ func hasUnresolvedVarRef(original, expanded string) bool {
|
||||
}
|
||||
|
||||
// expandWithEnv expands ${VAR} and $VAR references in s using the env map.
|
||||
// Falls back to the platform process env if a var isn't in the map.
|
||||
// Shell variables must start with a letter or '_' per POSIX; invalid identifiers
|
||||
// are returned literally so that "$100" and "$5" stay as-is.
|
||||
// Falls back to the platform process env only when the whole value is a
|
||||
// single variable reference; embedded process-env expansion is too broad for
|
||||
// imported org YAML because host variables such as HOME are not template data.
|
||||
func expandWithEnv(s string, env map[string]string) string {
|
||||
return os.Expand(s, func(key string) string {
|
||||
if len(key) == 0 {
|
||||
return "$"
|
||||
if s == "" {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
for i := 0; i < len(s); {
|
||||
if s[i] != '$' {
|
||||
b.WriteByte(s[i])
|
||||
i++
|
||||
continue
|
||||
}
|
||||
c := key[0]
|
||||
if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_') {
|
||||
return "$" + key // not a valid shell identifier — return literal
|
||||
|
||||
if i+1 >= len(s) {
|
||||
b.WriteByte('$')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if v, ok := env[key]; ok {
|
||||
return v
|
||||
|
||||
if s[i+1] == '{' {
|
||||
end := strings.IndexByte(s[i+2:], '}')
|
||||
if end < 0 {
|
||||
b.WriteByte('$')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
end += i + 2
|
||||
key := s[i+2 : end]
|
||||
ref := s[i : end+1]
|
||||
b.WriteString(expandEnvRef(key, ref, s, env))
|
||||
i = end + 1
|
||||
continue
|
||||
}
|
||||
return os.Getenv(key)
|
||||
})
|
||||
|
||||
if !isEnvIdentStart(s[i+1]) {
|
||||
b.WriteByte('$')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
j := i + 2
|
||||
for j < len(s) && isEnvIdentPart(s[j]) {
|
||||
j++
|
||||
}
|
||||
key := s[i+1 : j]
|
||||
ref := s[i:j]
|
||||
b.WriteString(expandEnvRef(key, ref, s, env))
|
||||
i = j
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// loadWorkspaceEnv reads the org root .env and the workspace-specific .env
|
||||
// expandEnvRef resolves a single variable reference extracted from s.
|
||||
//
|
||||
// Guards:
|
||||
// - Empty key → "$$" escape, return "$"
|
||||
// - key[0] not POSIX ident start → "$" + partial chars, return "$<chars>"
|
||||
// - Key in env map → return the mapped value (template override wins)
|
||||
// - Otherwise → only fall back to os.Getenv if the whole input string IS the
|
||||
// variable reference (ref == whole).
|
||||
//
|
||||
// Bare $VAR format:
|
||||
// $HOME (alone) → ref==whole → os.Getenv ✓ (host HOME is org-template HOME)
|
||||
// $HOME/path (partial) → ref!=whole → literal "$HOME" ✓ (CWE-78: prevents host leak)
|
||||
//
|
||||
// Braced ${VAR} format:
|
||||
// ${HOME} (alone) → ref==whole → os.Getenv ✓
|
||||
// ${ROLE}/admin (partial) → ref!=whole → literal ✓
|
||||
// "yes and ${NOT_SET}" (embedded) → ref!=whole → literal ✓
|
||||
//
|
||||
// This is the CWE-78 fix from commit a3a358f9.
|
||||
func expandEnvRef(key, ref, whole string, env map[string]string) string {
|
||||
if key == "" {
|
||||
return "$"
|
||||
}
|
||||
if !isEnvIdentStart(key[0]) {
|
||||
return "$" + key
|
||||
}
|
||||
if v, ok := env[key]; ok {
|
||||
return v
|
||||
}
|
||||
if ref == whole {
|
||||
return os.Getenv(key)
|
||||
}
|
||||
return ref
|
||||
}
|
||||
|
||||
func isEnvIdentStart(c byte) bool {
|
||||
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_'
|
||||
}
|
||||
|
||||
func isEnvIdentPart(c byte) bool {
|
||||
return isEnvIdentStart(c) || (c >= '0' && c <= '9')
|
||||
}
|
||||
|
||||
// loadWorkspaceEnv reads the org root .env and the workspace-specific .env .env and the workspace-specific .env
|
||||
// (workspace overrides org root). Used by both secret injection and channel
|
||||
// config expansion.
|
||||
//
|
||||
|
||||
@@ -462,8 +462,9 @@ func TestExpandWithEnv_LiteralDollar(t *testing.T) {
|
||||
func TestExpandWithEnv_PartiallyPresent(t *testing.T) {
|
||||
env := map[string]string{"SET": "yes"}
|
||||
result := expandWithEnv("${SET} and ${NOT_SET}", env)
|
||||
// ${SET} resolved; ${NOT_SET} -> "" via empty fallback.
|
||||
assert.Equal(t, "yes and ", result)
|
||||
// ${SET} resolved from env; ${NOT_SET} stays literal (not whole-string ref,
|
||||
// so os.Getenv fallback is NOT used — CWE-78 regression guard).
|
||||
assert.Equal(t, "yes and ${NOT_SET}", result)
|
||||
}
|
||||
|
||||
// mergeCategoryRouting tests — unions defaults with per-workspace routing.
|
||||
|
||||
@@ -276,3 +276,121 @@ func TestMergeCategoryRouting_OriginalMapsUnmodified(t *testing.T) {
|
||||
t.Error("ws routing should be unmodified after merge")
|
||||
}
|
||||
}
|
||||
|
||||
// ── expandWithEnv ─────────────────────────────────────────────────────────────
|
||||
//
|
||||
// CWE-78 regression tests. The original fix (a3a358f9) ensures that partial
|
||||
// variable references like $HOME/path are NOT resolved via os.Getenv — the
|
||||
// host HOME env var must not leak into org template values. Only whole-string
|
||||
// references ($VAR or ${VAR}) may fall back to the host process environment.
|
||||
|
||||
func TestExpandWithEnv_PartialRefDollarHomePath(t *testing.T) {
|
||||
// $HOME/path must NOT resolve to the host's HOME env var.
|
||||
// The literal $HOME must be returned as-is.
|
||||
got := expandWithEnv("$HOME/path", nil)
|
||||
if got != "$HOME/path" {
|
||||
t.Errorf("$HOME/path: got %q, want literal $HOME/path", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_PartialRefBracedRoleAdmin(t *testing.T) {
|
||||
// ${ROLE}/admin — ROLE is not in env, so expand to the literal ${ROLE}/admin.
|
||||
got := expandWithEnv("${ROLE}/admin", nil)
|
||||
if got != "${ROLE}/admin" {
|
||||
t.Errorf("${ROLE}/admin: got %q, want literal ${ROLE}/admin", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_PartialRefMiddleOfString(t *testing.T) {
|
||||
// $ROLE in the middle of a string — literal, not os.Getenv.
|
||||
got := expandWithEnv("prefix/$ROLE/suffix", nil)
|
||||
if got != "prefix/$ROLE/suffix" {
|
||||
t.Errorf("prefix/$ROLE/suffix: got %q, want literal", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_WholeVarInEnv(t *testing.T) {
|
||||
// Whole-string $VAR that IS in env — env value wins.
|
||||
env := map[string]string{"FOO": "barvalue"}
|
||||
got := expandWithEnv("$FOO", env)
|
||||
if got != "barvalue" {
|
||||
t.Errorf("$FOO with FOO=barvalue: got %q, want barvalue", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_WholeVarBracedInEnv(t *testing.T) {
|
||||
// Whole-string ${VAR} that IS in env — env value wins.
|
||||
env := map[string]string{"FOO": "barvalue"}
|
||||
got := expandWithEnv("${FOO}", env)
|
||||
if got != "barvalue" {
|
||||
t.Errorf("${FOO} with FOO=barvalue: got %q, want barvalue", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_WholeVarNotInEnvBare(t *testing.T) {
|
||||
// Whole-string $VAR not in env — falls back to os.Getenv.
|
||||
// If the host has the var, we get the host value. If not, empty.
|
||||
// At minimum, the result must NOT be the literal "$UNDEFINED_VAR_9Z".
|
||||
got := expandWithEnv("$UNDEFINED_VAR_9Z", nil)
|
||||
if got == "$UNDEFINED_VAR_9Z" {
|
||||
t.Errorf("$UNDEFINED_VAR_9Z: should expand (whole-string fallback to os.Getenv), got literal")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_WholeVarNotInEnvBraced(t *testing.T) {
|
||||
// Whole-string ${VAR} not in env — falls back to os.Getenv.
|
||||
got := expandWithEnv("${UNDEFINED_VAR_9Z}", nil)
|
||||
if got == "${UNDEFINED_VAR_9Z}" {
|
||||
t.Errorf("${UNDEFINED_VAR_9Z}: should expand (whole-string fallback to os.Getenv), got literal")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_EmptyString(t *testing.T) {
|
||||
got := expandWithEnv("", map[string]string{"FOO": "bar"})
|
||||
if got != "" {
|
||||
t.Errorf("empty string: got %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_NoVarRefs(t *testing.T) {
|
||||
got := expandWithEnv("plain string with no vars", map[string]string{"FOO": "bar"})
|
||||
if got != "plain string with no vars" {
|
||||
t.Errorf("plain string: got %q, want unchanged", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_MultipleVarRefs(t *testing.T) {
|
||||
// Two vars, both whole — both expand from env.
|
||||
env := map[string]string{"A": "alpha", "B": "beta"}
|
||||
got := expandWithEnv("$A and $B and more", env)
|
||||
if got != "alpha and beta and more" {
|
||||
t.Errorf("multiple vars: got %q, want alpha and beta and more", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_NumericVarRef(t *testing.T) {
|
||||
// $5 — starts with digit, not a valid identifier start.
|
||||
// Must return the literal "$5", not expand via os.Getenv.
|
||||
got := expandWithEnv("$5", map[string]string{"5": "five"})
|
||||
if got != "$5" {
|
||||
t.Errorf("$5: got %q, want literal $5", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_DollarEscape(t *testing.T) {
|
||||
// $$ → both $ written literally (each $ is not followed by an identifier char,
|
||||
// so it is written as-is). No special escape sequence for $$.
|
||||
got := expandWithEnv("$$", nil)
|
||||
if got != "$$" {
|
||||
t.Errorf("$$: got %q, want literal $$", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandWithEnv_MixedPartialAndWhole(t *testing.T) {
|
||||
// $A is in env (whole), $HOME is partial — only $A expands.
|
||||
env := map[string]string{"A": "alpha"}
|
||||
got := expandWithEnv("$A at $HOME", env)
|
||||
if got != "alpha at $HOME" {
|
||||
t.Errorf("$A at $HOME: got %q, want alpha at $HOME", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,12 +4,14 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -156,6 +158,10 @@ type cpProvisionRequest struct {
|
||||
Tier int `json:"tier"`
|
||||
PlatformURL string `json:"platform_url"`
|
||||
Env map[string]string `json:"env"`
|
||||
// ConfigFiles are base64-encoded config files collected from the template
|
||||
// directory and inline ConfigFiles. Sent to the control plane so CP can
|
||||
// write them into the EC2 instance at boot time (OFFSEC-010 wiring fix).
|
||||
ConfigFiles map[string]string `json:"config_files,omitempty"`
|
||||
}
|
||||
|
||||
type cpProvisionResponse struct {
|
||||
@@ -179,6 +185,16 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
|
||||
}
|
||||
env["ADMIN_TOKEN"] = p.adminToken
|
||||
}
|
||||
|
||||
// Collect config files from the template directory and inline ConfigFiles.
|
||||
// Without this wiring, OFFSEC-010 config-file collection is dead code and
|
||||
// CP receives no files. An error here is fatal — a workspace with no
|
||||
// config files cannot function.
|
||||
configFiles, err := collectCPConfigFiles(cfg)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("cp provisioner: collect config files: %w", err)
|
||||
}
|
||||
|
||||
req := cpProvisionRequest{
|
||||
OrgID: p.orgID,
|
||||
WorkspaceID: cfg.WorkspaceID,
|
||||
@@ -186,6 +202,7 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
|
||||
Tier: cfg.Tier,
|
||||
PlatformURL: cfg.PlatformURL,
|
||||
Env: env,
|
||||
ConfigFiles: configFiles,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(req)
|
||||
@@ -237,6 +254,92 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
|
||||
return result.InstanceID, nil
|
||||
}
|
||||
|
||||
const cpConfigFilesMaxBytes = 12 << 10
|
||||
|
||||
// collectCPConfigFiles walks cfg.TemplatePath and collects regular files as
|
||||
// base64-encoded key→value entries, plus any inline cfg.ConfigFiles. Used to
|
||||
// ship template config into the EC2 instance via the control plane.
|
||||
//
|
||||
// Security (OFFSEC-010):
|
||||
// - cfg.TemplatePath itself must not be a symlink (os.Lstat check).
|
||||
// - WalkDir skips symlinks inside the tree (d.Type()&os.ModeSymlink guard).
|
||||
// - All file paths are cleaned and relativized; absolute/traversal paths
|
||||
// are rejected by addFile's sanitization checks.
|
||||
//
|
||||
// Returns nil, nil when no files are found — caller must handle empty case.
|
||||
func collectCPConfigFiles(cfg WorkspaceConfig) (map[string]string, error) {
|
||||
files := make(map[string]string)
|
||||
total := 0
|
||||
addFile := func(name string, data []byte) error {
|
||||
name = filepath.ToSlash(filepath.Clean(name))
|
||||
if name == "." || strings.HasPrefix(name, "../") || strings.HasPrefix(name, "/") || strings.Contains(name, "/../") {
|
||||
return fmt.Errorf("invalid config file path %q", name)
|
||||
}
|
||||
total += len(data)
|
||||
if total > cpConfigFilesMaxBytes {
|
||||
return fmt.Errorf("config files exceed %d bytes", cpConfigFilesMaxBytes)
|
||||
}
|
||||
files[name] = base64.StdEncoding.EncodeToString(data)
|
||||
return nil
|
||||
}
|
||||
|
||||
if cfg.TemplatePath != "" {
|
||||
// Reject symlinks on the root itself — WalkDir follows symlinks,
|
||||
// so a symlink TemplatePath that escapes the intended root directory
|
||||
// would bypass the subsequent path-relativization checks below.
|
||||
rootInfo, err := os.Lstat(cfg.TemplatePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("collectCPConfigFiles: lstat template path: %w", err)
|
||||
}
|
||||
if rootInfo.Mode()&os.ModeSymlink != 0 {
|
||||
return nil, fmt.Errorf("collectCPConfigFiles: template path must not be a symlink")
|
||||
}
|
||||
err = filepath.WalkDir(cfg.TemplatePath, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
}
|
||||
// Skip symlinks — WalkDir follows them by default, which means
|
||||
// a symlink inside the template dir pointing to /etc/passwd
|
||||
// would be traversed even though the resulting relative-path
|
||||
// check would correctly reject it. Defense-in-depth. (OFFSEC-010)
|
||||
if d.Type()&os.ModeSymlink != 0 {
|
||||
return nil
|
||||
}
|
||||
if d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !info.Mode().IsRegular() {
|
||||
return nil
|
||||
}
|
||||
rel, err := filepath.Rel(cfg.TemplatePath, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return addFile(rel, data)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
for name, data := range cfg.ConfigFiles {
|
||||
if err := addFile(name, data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if len(files) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// Stop terminates the workspace's EC2 instance via the control plane.
|
||||
//
|
||||
// Looks up the actual EC2 instance_id from the workspaces table before
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -842,3 +844,167 @@ func TestIsRunning_EmptyInstanceIDReturnsFalse(t *testing.T) {
|
||||
t.Errorf("IsRunning with empty instance_id should return running=false, got true")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCollectCPConfigFiles_SkipsSymlinks — WalkDir follows symlinks by default,
|
||||
// but collectCPConfigFiles must skip them so a symlink inside a template dir
|
||||
// pointing outside (e.g. ln -s /etc snapshot) cannot be traversed.
|
||||
// Verifies OFFSEC-010 defense-in-depth fix.
|
||||
func TestCollectCPConfigFiles_SkipsSymlinks(t *testing.T) {
|
||||
tmpl := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(tmpl, "config.yaml"), []byte("name: real\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sensitiveDir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(sensitiveDir, "secret.txt"), []byte("SENSITIVE\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
symlinkPath := filepath.Join(tmpl, "snapshot")
|
||||
if err := os.Symlink(sensitiveDir, symlinkPath); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
files, err := collectCPConfigFiles(WorkspaceConfig{TemplatePath: tmpl})
|
||||
if err != nil {
|
||||
t.Fatalf("collectCPConfigFiles: %v", err)
|
||||
}
|
||||
if files == nil {
|
||||
t.Fatal("files should not be nil")
|
||||
}
|
||||
if _, ok := files["config.yaml"]; !ok {
|
||||
t.Errorf("config.yaml missing from files")
|
||||
}
|
||||
for k := range files {
|
||||
if strings.Contains(k, "snapshot") || strings.Contains(k, "secret") {
|
||||
t.Errorf("symlink path %q should not be in files — OFFSEC-010 regression", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCollectCPConfigFiles_RejectsRootSymlink — if cfg.TemplatePath itself is
|
||||
// a symlink, WalkDir would follow it to an arbitrary directory, bypassing the
|
||||
// cfg.TemplatePath boundary. The function must reject this explicitly.
|
||||
func TestCollectCPConfigFiles_RejectsRootSymlink(t *testing.T) {
|
||||
real := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(real, "config.yaml"), []byte("name: real\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
link := filepath.Join(t.TempDir(), "template-link")
|
||||
if err := os.Symlink(real, link); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err := collectCPConfigFiles(WorkspaceConfig{TemplatePath: link})
|
||||
if err == nil {
|
||||
t.Error("collectCPConfigFiles with symlink TemplatePath should return error")
|
||||
}
|
||||
if err != nil && !strings.Contains(err.Error(), "symlink") {
|
||||
t.Errorf("expected symlink-related error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_PassesConfigFilesToCP — Start must call collectCPConfigFiles and
|
||||
// include the result (base64-encoded files) in the cpProvisionRequest body sent
|
||||
// to the control plane. Without this wiring, OFFSEC-010 config-file collection
|
||||
// is dead code and CP receives no files. (OFFSEC-010 wiring fix)
|
||||
func TestStart_PassesConfigFilesToCP(t *testing.T) {
|
||||
tmpl := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(tmpl, "agent.yaml"), []byte("name: test\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inlineFiles := map[string][]byte{
|
||||
"runtime.yaml": []byte("runtime: python\n"),
|
||||
}
|
||||
|
||||
var sawBody cpProvisionRequest
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&sawBody); err != nil {
|
||||
t.Fatalf("decode CP request body: %v", err)
|
||||
}
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
_, _ = io.WriteString(w, `{"instance_id":"i-cfgt","state":"running"}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &CPProvisioner{baseURL: srv.URL, orgID: "org-1", httpClient: srv.Client()}
|
||||
_, err := p.Start(context.Background(), WorkspaceConfig{
|
||||
WorkspaceID: "ws-cfg",
|
||||
Runtime: "python",
|
||||
Tier: 2,
|
||||
PlatformURL: "https://tenant.example",
|
||||
TemplatePath: tmpl,
|
||||
ConfigFiles: inlineFiles,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
if sawBody.WorkspaceID != "ws-cfg" {
|
||||
t.Errorf("WorkspaceID = %q, want ws-cfg", sawBody.WorkspaceID)
|
||||
}
|
||||
if sawBody.ConfigFiles == nil {
|
||||
t.Fatal("ConfigFiles must not be nil — collectCPConfigFiles returned nil")
|
||||
}
|
||||
foundAgent := false
|
||||
foundRuntime := false
|
||||
for name := range sawBody.ConfigFiles {
|
||||
if name == "agent.yaml" {
|
||||
foundAgent = true
|
||||
}
|
||||
if name == "runtime.yaml" {
|
||||
foundRuntime = true
|
||||
}
|
||||
}
|
||||
if !foundAgent {
|
||||
t.Errorf("agent.yaml missing from ConfigFiles; keys = %v", mapKeys(sawBody.ConfigFiles))
|
||||
}
|
||||
if !foundRuntime {
|
||||
t.Errorf("runtime.yaml missing from ConfigFiles; keys = %v", mapKeys(sawBody.ConfigFiles))
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_CollectConfigFilesErrorSurfaces — if collectCPConfigFiles returns
|
||||
// an error (e.g. TemplatePath is a symlink, or files exceed the 12 KiB cap),
|
||||
// Start must propagate that error rather than silently continuing. (OFFSEC-010)
|
||||
func TestStart_CollectConfigFilesErrorSurfaces(t *testing.T) {
|
||||
realDir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(realDir, "config.yaml"), []byte("data\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
symlinkPath := filepath.Join(t.TempDir(), "template-link")
|
||||
if err := os.Symlink(realDir, symlinkPath); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Error("CP must not be called when collectCPConfigFiles fails")
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &CPProvisioner{baseURL: srv.URL, orgID: "org-1", httpClient: srv.Client()}
|
||||
_, err := p.Start(context.Background(), WorkspaceConfig{
|
||||
WorkspaceID: "ws-symlink",
|
||||
Runtime: "python",
|
||||
TemplatePath: symlinkPath,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error when TemplatePath is a symlink, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "symlink") {
|
||||
t.Errorf("error should mention symlink, got: %v", err)
|
||||
}
|
||||
if strings.Contains(err.Error(), "cp provisioner: send") {
|
||||
t.Errorf("CP must not be called; error should be from collectConfigFiles, not HTTP send: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// mapKeys returns the keys of a map — helper for test assertions.
|
||||
func mapKeys(m map[string]string) []string {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
@@ -40,6 +40,8 @@ _A2A_BOUNDARY_END = "[/A2A_RESULT_FROM_PEER]"
|
||||
# inside the trusted zone. Escape BOTH boundary markers in the raw text
|
||||
# before wrapping so they can never close the boundary early.
|
||||
# We use "[/ " as the escape prefix — visually distinct from the real marker.
|
||||
_A2A_BOUNDARY_START_ESCAPED = "[/ A2A_RESULT_FROM_PEER]"
|
||||
_A2A_BOUNDARY_END_ESCAPED = "[/ /A2A_RESULT_FROM_PEER]"
|
||||
|
||||
|
||||
def _escape_boundary_markers(text: str) -> str:
|
||||
@@ -50,8 +52,8 @@ def _escape_boundary_markers(text: str) -> str:
|
||||
the boundary early or inject a fake opener.
|
||||
"""
|
||||
return (
|
||||
text.replace(_A2A_BOUNDARY_START, "[/ A2A_RESULT_FROM_PEER]")
|
||||
.replace(_A2A_BOUNDARY_END, "[/ /A2A_RESULT_FROM_PEER]")
|
||||
text.replace(_A2A_BOUNDARY_START, _A2A_BOUNDARY_START_ESCAPED)
|
||||
.replace(_A2A_BOUNDARY_END, _A2A_BOUNDARY_END_ESCAPED)
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -49,7 +49,9 @@ from a2a_client import (
|
||||
from a2a_tools_rbac import auth_headers_for_heartbeat as _auth_headers_for_heartbeat
|
||||
from _sanitize_a2a import (
|
||||
_A2A_BOUNDARY_END,
|
||||
_A2A_BOUNDARY_END_ESCAPED,
|
||||
_A2A_BOUNDARY_START,
|
||||
_A2A_BOUNDARY_START_ESCAPED,
|
||||
sanitize_a2a_result,
|
||||
) # noqa: E402
|
||||
|
||||
@@ -330,8 +332,18 @@ async def tool_delegate_task(
|
||||
# markers so the agent can distinguish trusted (own output) from untrusted
|
||||
# (peer-supplied) content. Explicit wrapping here rather than inside
|
||||
# sanitize_a2a_result preserves a clean separation of concerns.
|
||||
#
|
||||
# Truncate at the closer BEFORE sanitizing so the raw closer (which gets
|
||||
# lost during escaping) is removed from the content. After truncation,
|
||||
# sanitize the remaining text and wrap with escaped boundary markers.
|
||||
if _A2A_BOUNDARY_END in result:
|
||||
result = result[:result.index(_A2A_BOUNDARY_END)]
|
||||
escaped = sanitize_a2a_result(result)
|
||||
return f"{_A2A_BOUNDARY_START}\n{escaped}\n{_A2A_BOUNDARY_END}"
|
||||
return (
|
||||
f"{_A2A_BOUNDARY_START_ESCAPED}\n"
|
||||
f"{escaped}\n"
|
||||
f"{_A2A_BOUNDARY_END_ESCAPED}"
|
||||
)
|
||||
|
||||
|
||||
async def tool_delegate_task_async(
|
||||
|
||||
@@ -218,7 +218,8 @@ class TestPollingPathSanitization:
|
||||
result = asyncio.run(d.tool_delegate_task("ws-peer", "do it"))
|
||||
# tool_delegate_task wraps the sanitized text in _A2A_BOUNDARY_START/END
|
||||
# (NOT _A2A_RESULT_FROM_PEER — that marker is for the messaging path).
|
||||
assert d._A2A_BOUNDARY_START in result
|
||||
assert d._A2A_BOUNDARY_END in result
|
||||
# Wrapped in escaped form to prevent raw closer from appearing in output.
|
||||
assert d._A2A_BOUNDARY_START_ESCAPED in result
|
||||
assert d._A2A_BOUNDARY_END_ESCAPED in result
|
||||
assert "Sanitized peer reply" in result
|
||||
|
||||
|
||||
@@ -277,7 +277,7 @@ class TestToolDelegateTask:
|
||||
patch("a2a_tools.report_activity", new=AsyncMock()):
|
||||
result = await a2a_tools.tool_delegate_task("ws-1", "do something")
|
||||
|
||||
assert result == "[A2A_RESULT_FROM_PEER]\nTask completed!\n[/A2A_RESULT_FROM_PEER]"
|
||||
assert result == "[/ A2A_RESULT_FROM_PEER]\nTask completed!\n[/ /A2A_RESULT_FROM_PEER]"
|
||||
|
||||
async def test_error_response_returns_delegation_failed_message(self):
|
||||
"""When send_a2a_message returns _A2A_ERROR_PREFIX text, delegation fails."""
|
||||
@@ -305,7 +305,7 @@ class TestToolDelegateTask:
|
||||
patch("a2a_tools.report_activity", new=AsyncMock()):
|
||||
result = await a2a_tools.tool_delegate_task("ws-cached", "task")
|
||||
|
||||
assert result == "[A2A_RESULT_FROM_PEER]\ndone\n[/A2A_RESULT_FROM_PEER]"
|
||||
assert result == "[/ A2A_RESULT_FROM_PEER]\ndone\n[/ /A2A_RESULT_FROM_PEER]"
|
||||
|
||||
async def test_peer_name_falls_back_to_id_prefix(self):
|
||||
"""When peer has no name and cache is empty, name = first 8 chars of workspace_id."""
|
||||
@@ -319,7 +319,7 @@ class TestToolDelegateTask:
|
||||
patch("a2a_tools.report_activity", new=AsyncMock()):
|
||||
result = await a2a_tools.tool_delegate_task("ws-nona000", "task")
|
||||
|
||||
assert result == "[A2A_RESULT_FROM_PEER]\nok\n[/A2A_RESULT_FROM_PEER]"
|
||||
assert result == "[/ A2A_RESULT_FROM_PEER]\nok\n[/ /A2A_RESULT_FROM_PEER]"
|
||||
# Cache should now have been set
|
||||
assert a2a_tools._peer_names.get("ws-nona000") is not None
|
||||
|
||||
|
||||
@@ -69,7 +69,7 @@ class TestFlagOffLegacyPath:
|
||||
monkeypatch.delenv("DELEGATION_SYNC_VIA_INBOX", raising=False)
|
||||
|
||||
import a2a_tools
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END, _A2A_BOUNDARY_START
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END_ESCAPED, _A2A_BOUNDARY_START_ESCAPED
|
||||
send_calls = []
|
||||
|
||||
async def fake_send(workspace_id, task, source_workspace_id=None):
|
||||
@@ -91,8 +91,8 @@ class TestFlagOffLegacyPath:
|
||||
)
|
||||
|
||||
# OFFSEC-003: result is wrapped in boundary markers
|
||||
assert _A2A_BOUNDARY_START in result
|
||||
assert _A2A_BOUNDARY_END in result
|
||||
assert _A2A_BOUNDARY_START_ESCAPED in result
|
||||
assert _A2A_BOUNDARY_END_ESCAPED in result
|
||||
assert "legacy ok" in result
|
||||
assert send_calls == [("ws-target", "task body", "ws-self")]
|
||||
poll_mock.assert_not_called()
|
||||
@@ -124,7 +124,7 @@ class TestPollModeAutoFallback:
|
||||
monkeypatch.delenv("DELEGATION_SYNC_VIA_INBOX", raising=False)
|
||||
|
||||
import a2a_tools
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END, _A2A_BOUNDARY_START
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END_ESCAPED, _A2A_BOUNDARY_START_ESCAPED
|
||||
from a2a_client import _A2A_QUEUED_PREFIX
|
||||
|
||||
send_calls = []
|
||||
@@ -159,8 +159,8 @@ class TestPollModeAutoFallback:
|
||||
assert poll_calls[0] == ("ws-target", "task body", "ws-self")
|
||||
# Caller sees the real reply, NOT the queued sentinel and NOT
|
||||
# a DELEGATION FAILED string. Wrapped in OFFSEC-003 boundary markers.
|
||||
assert _A2A_BOUNDARY_START in result
|
||||
assert _A2A_BOUNDARY_END in result
|
||||
assert _A2A_BOUNDARY_START_ESCAPED in result
|
||||
assert _A2A_BOUNDARY_END_ESCAPED in result
|
||||
assert "real response from poll-mode peer" in result
|
||||
|
||||
async def test_non_queued_send_result_does_not_trigger_fallback(self, monkeypatch):
|
||||
@@ -169,7 +169,7 @@ class TestPollModeAutoFallback:
|
||||
monkeypatch.delenv("DELEGATION_SYNC_VIA_INBOX", raising=False)
|
||||
|
||||
import a2a_tools
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END, _A2A_BOUNDARY_START
|
||||
from _sanitize_a2a import _A2A_BOUNDARY_END_ESCAPED, _A2A_BOUNDARY_START_ESCAPED
|
||||
|
||||
async def fake_send(*_a, **_kw):
|
||||
return "normal reply"
|
||||
@@ -189,8 +189,8 @@ class TestPollModeAutoFallback:
|
||||
)
|
||||
|
||||
# OFFSEC-003: wrapped in boundary markers
|
||||
assert _A2A_BOUNDARY_START in result
|
||||
assert _A2A_BOUNDARY_END in result
|
||||
assert _A2A_BOUNDARY_START_ESCAPED in result
|
||||
assert _A2A_BOUNDARY_END_ESCAPED in result
|
||||
assert "normal reply" in result
|
||||
poll_mock.assert_not_called()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user