mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-05-19 16:28:28 +00:00
feat(acp): LLM-based message rewrite middleware with custom prompts (#3191)
* feat(acp): LLM-based message rewrite middleware
Add MessageRewriteMiddleware that intercepts ACP messages and appends
LLM-rewritten versions with _meta.rewritten=true at turn boundaries.
Original messages pass through unmodified. At the end of each turn
(before tool calls or at response end), accumulated thought/message
chunks are sent to LLM for rewriting into business-friendly text.
- TurnBuffer: accumulates chunks per turn
- LlmRewriter: calls LLM with configurable prompt
- MessageRewriteMiddleware: orchestrates intercept → buffer → rewrite → emit
- BaseEmitter.sendUpdate: routes through middleware when configured
- Session: initializes middleware from settings.messageRewrite config
Enable via settings.json:
{
"messageRewrite": {
"enabled": true,
"target": "both",
"prompt": "custom system prompt for rewriter"
}
}
Rewritten messages carry _meta.rewritten=true for frontend to
prioritize display. Original messages remain for debugging.
* fix: TypeScript 编译错误修复 + 优化默认改写 prompt(参考竞品风格)
* fix: 从 user/workspace originalSettings 读取 messageRewrite 配置(绕过 schema 校验)
* feat: 非交互 CLI 模式也支持 message rewrite(eval 可用)
* fix: 禁用 rewriter LLM 的 thinking,过滤 thought 部分只取纯文本输出
* fix: cron 路径补齐 message rewrite flush + 代码质量优化
- Session.ts cron 路径添加 messageRewriter.flushTurn() 调用
- nonInteractiveCli.ts cron 路径添加 turnBuffer 累积 + flush + rewrite
- 提取 loadRewriteConfig() 共享函数,消除两处重复配置读取
- 主路径和 cron 路径添加 turnBuffer.markToolCall()
- rewrite 调用添加 30s 超时保护(AbortSignal.timeout)
- 修复 import 语句被 const 声明分割的问题
* feat: rewrite 支持 async/sync 模式(默认 async,不增加执行时间)
* feat: rewrite prompt 通用化 + 上下文连贯 + promptFile + async 修复
- 默认 prompt 改为通用英文版(适配任意 coding agent,不绑定数据分析场景)
- 支持 promptFile 配置项,从文件加载自定义 prompt(优先于 inline prompt)
- 上下文连贯性:lastOutput 记录上一轮改写结果,拼接到下一轮输入,
避免连续 turn 间信息重复
- 修复 CLI 非交互模式 async rewrite 丢失:void doRewrite() 改为
pendingRewrites 数组 + emitResult 前 Promise.allSettled
- 增加 debug logging:REWRITE INPUT/OUTPUT 完整内容 + prev_output 长度
* refactor: remove sync rewrite mode, always use async (non-blocking) rewrite
- Remove `async` field from MessageRewriteConfig
- MessageRewriteMiddleware.flushTurn() always fires in background
- nonInteractiveCli.ts main & cron paths always push to pendingRewrites
- No user-facing latency from rewrite calls
* fix: address review feedback — trust check, timeout, history replay
1. loadRewriteConfig: skip workspace settings when !isTrusted, preventing
untrusted repos from enabling rewriter with a custom prompt
2. MessageRewriteMiddleware.flushTurn: always enforce 30s timeout internally,
even when caller provides no AbortSignal (interactive path)
3. Install rewriter AFTER history replay completes (Session.installRewriter),
so historical messages are never rewritten on session load
* fix: address second round review — target filter, timeout, rewrite queue
1. nonInteractiveCli: apply rewriteConfig.target filter to accumulation
(main path and cron path), matching MessageRewriteMiddleware behavior
2. nonInteractiveCli: add 30s AbortSignal.timeout to rewrite calls in
both main and cron paths
3. MessageRewriteMiddleware: replace single pendingRewrite slot with
pendingRewrites array + Promise.allSettled, ensuring all rewrites
complete before session exits
* test: add unit tests for TurnBuffer, loadRewriteConfig, MessageRewriteMiddleware
- TurnBuffer: flush, reset, isEmpty, markToolCall, whitespace filtering (12 tests)
- loadRewriteConfig: isTrusted gating, workspace/user precedence (5 tests)
- MessageRewriteMiddleware: target filtering, tool_call boundary flush,
pendingRewrites queue, rewrite metadata (9 tests)
* fix: config.test.ts use unknown cast for LoadedSettings stub (fix tsc --build)
* fix: filter LLM literal "empty string" responses in rewriter output
LLM sometimes outputs "(空字符串)" or similar text instead of actual
empty string when instructed to "return empty string". Add regex patterns
to catch common variants and treat them as null (skip rewrite output).
* revert: remove LLM empty-string pattern defense, rely on prompt fix instead
* fix: prevent async rewrite from corrupting adapter state + honor config.model
1. nonInteractiveCli: rewrite promises now return data only, adapter
emission happens synchronously via emitSettledRewrites() at safe
boundaries (before next turn starts, before cron next turn, before
final result). Prevents concurrent startAssistantMessage corruption.
2. LlmRewriter: use rewriteConfig.model when set, fallback to
config.getModel(). Previously model field was defined but ignored.
* docs: add messageRewrite configuration guide to settings.md
* Revert "docs: add messageRewrite configuration guide to settings.md"
This reverts commit ecd57e2d5a.
* feat: add contextTurns config for rewrite history context
Allow configuring how many previous rewrite outputs are included as
context when rewriting a new turn:
- contextTurns: 1 (default) = last rewrite only
- contextTurns: 0 = no context
- contextTurns: N = last N rewrites
- contextTurns: "all" = all previous rewrites
* refactor: rename target 'both' to 'all' + add LlmRewriter unit tests
- Rename target value 'both' → 'all' for future extensibility (e.g. 'tool')
- Add LlmRewriter tests: contextTurns (0/1/N/all), model override, filtering
- Total: 35 tests across 4 test files
* refactor: remove message rewrite from non-interactive CLI mode
Non-interactive mode (qwen -p "..." --output-format json) consumers are
scripts/programs that don't need user-friendly rewrites. Additionally,
the JSON output adapter doesn't support _meta fields, so rewritten text
was silently mixed into normal assistant messages without any marker.
Rewrite middleware is now ACP-only (Session path).
* revert: restore package-lock.json and nonInteractiveCli.ts to main state
* docs: add README for message rewrite middleware
Explain the feature purpose (business-oriented output customization),
mark it as a temporary solution, and reference the hook-based
alternative (#3266) for future discussion.
* docs: move temporary-solution notice to top of README
* docs: simplify temporary-solution notice in rewrite README
This commit is contained in:
parent
3c556c01f3
commit
ae424e004b
15 changed files with 1180 additions and 0 deletions
|
|
@ -526,6 +526,9 @@ class QwenAgent implements Agent {
|
|||
await session.replayHistory(conversation.messages);
|
||||
}
|
||||
|
||||
// Install rewriter AFTER history replay to avoid rewriting historical messages
|
||||
session.installRewriter();
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -88,6 +88,10 @@ import {
|
|||
buildPermissionRequestContent,
|
||||
toPermissionOptions,
|
||||
} from './permissionUtils.js';
|
||||
import {
|
||||
MessageRewriteMiddleware,
|
||||
loadRewriteConfig,
|
||||
} from './rewrite/index.js';
|
||||
|
||||
const debugLogger = createDebugLogger('SESSION');
|
||||
|
||||
|
|
@ -124,6 +128,9 @@ export class Session implements SessionContext {
|
|||
private readonly planEmitter: PlanEmitter;
|
||||
private readonly messageEmitter: MessageEmitter;
|
||||
|
||||
// Message rewrite middleware (optional, installed after history replay)
|
||||
messageRewriter?: MessageRewriteMiddleware;
|
||||
|
||||
// Implement SessionContext interface
|
||||
readonly sessionId: string;
|
||||
|
||||
|
|
@ -152,6 +159,22 @@ export class Session implements SessionContext {
|
|||
return this.config;
|
||||
}
|
||||
|
||||
/**
|
||||
* Install the message rewrite middleware if configured.
|
||||
* Must be called AFTER history replay to avoid rewriting historical messages.
|
||||
*/
|
||||
installRewriter(): void {
|
||||
const rewriteConfig = loadRewriteConfig(this.settings);
|
||||
if (rewriteConfig?.enabled) {
|
||||
debugLogger.info('Message rewrite middleware enabled');
|
||||
this.messageRewriter = new MessageRewriteMiddleware(
|
||||
this.config,
|
||||
rewriteConfig,
|
||||
(update) => this.sendUpdate(update),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Replays conversation history to the client using modular components.
|
||||
* Delegates to HistoryReplayer for consistent event emission.
|
||||
|
|
@ -391,6 +414,11 @@ export class Session implements SessionContext {
|
|||
}
|
||||
|
||||
if (usageMetadata) {
|
||||
// Kick off rewrite in background (non-blocking, runs parallel to tools)
|
||||
if (this.messageRewriter) {
|
||||
this.messageRewriter.flushTurn(pendingSend.signal);
|
||||
}
|
||||
|
||||
const durationMs = Date.now() - streamStartTime;
|
||||
await this.messageEmitter.emitUsageMetadata(
|
||||
usageMetadata,
|
||||
|
|
@ -414,6 +442,10 @@ export class Session implements SessionContext {
|
|||
nextMessage = { role: 'user', parts: toolResponseParts };
|
||||
}
|
||||
}
|
||||
// Wait for any pending rewrite before returning
|
||||
if (this.messageRewriter) {
|
||||
await this.messageRewriter.waitForPendingRewrites();
|
||||
}
|
||||
return { stopReason: 'end_turn' };
|
||||
},
|
||||
);
|
||||
|
|
@ -560,6 +592,10 @@ export class Session implements SessionContext {
|
|||
}
|
||||
|
||||
if (usageMetadata) {
|
||||
// Kick off rewrite in background (non-blocking)
|
||||
if (this.messageRewriter) {
|
||||
this.messageRewriter.flushTurn(ac.signal);
|
||||
}
|
||||
const durationMs = Date.now() - streamStartTime;
|
||||
await this.messageEmitter.emitUsageMetadata(
|
||||
usageMetadata,
|
||||
|
|
|
|||
|
|
@ -31,8 +31,13 @@ export abstract class BaseEmitter {
|
|||
|
||||
/**
|
||||
* Sends a session update to the ACP client.
|
||||
* If a message rewriter is configured, updates pass through it first
|
||||
* (original messages are sent as-is, rewritten versions are appended).
|
||||
*/
|
||||
protected async sendUpdate(update: SessionUpdate): Promise<void> {
|
||||
if (this.ctx.messageRewriter) {
|
||||
return this.ctx.messageRewriter.interceptUpdate(update);
|
||||
}
|
||||
return this.ctx.sendUpdate(update);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,227 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import type { TurnContent, MessageRewriteConfig } from './types.js';
|
||||
|
||||
// Mock core to avoid Vite https resolution issue
|
||||
vi.mock('@qwen-code/qwen-code-core', () => ({
|
||||
createDebugLogger: () => ({
|
||||
info: vi.fn(),
|
||||
warn: vi.fn(),
|
||||
debug: vi.fn(),
|
||||
error: vi.fn(),
|
||||
}),
|
||||
}));
|
||||
|
||||
// Track generateContent calls
|
||||
const mockGenerateContent = vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [{ text: 'rewritten output' }],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const { LlmRewriter } = await import('./LlmRewriter.js');
|
||||
|
||||
function makeConfig(): Config {
|
||||
return {
|
||||
getContentGenerator: () => ({
|
||||
generateContent: mockGenerateContent,
|
||||
}),
|
||||
getModel: () => 'test-model',
|
||||
} as unknown as Config;
|
||||
}
|
||||
|
||||
function makeTurn(messages: string[], thoughts: string[] = []): TurnContent {
|
||||
return { messages, thoughts, hasToolCalls: false };
|
||||
}
|
||||
|
||||
describe('LlmRewriter', () => {
|
||||
beforeEach(() => {
|
||||
mockGenerateContent.mockClear();
|
||||
mockGenerateContent.mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'rewritten output' }] } }],
|
||||
});
|
||||
});
|
||||
|
||||
describe('contextTurns', () => {
|
||||
it('should include last rewrite output by default (contextTurns=1)', async () => {
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
// First call — no context
|
||||
await rewriter.rewrite(makeTurn(['first message']));
|
||||
const firstInput =
|
||||
mockGenerateContent.mock.calls[0][0].contents[0].parts[0].text;
|
||||
expect(firstInput).not.toContain('上一轮改写结果');
|
||||
|
||||
// Second call — should include first rewrite output
|
||||
await rewriter.rewrite(makeTurn(['second message']));
|
||||
const secondInput =
|
||||
mockGenerateContent.mock.calls[1][0].contents[0].parts[0].text;
|
||||
expect(secondInput).toContain('上一轮改写结果');
|
||||
expect(secondInput).toContain('rewritten output');
|
||||
});
|
||||
|
||||
it('should include no context when contextTurns=0', async () => {
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
contextTurns: 0,
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['first']));
|
||||
await rewriter.rewrite(makeTurn(['second']));
|
||||
|
||||
const secondInput =
|
||||
mockGenerateContent.mock.calls[1][0].contents[0].parts[0].text;
|
||||
expect(secondInput).not.toContain('上一轮改写结果');
|
||||
});
|
||||
|
||||
it('should include last N rewrites when contextTurns=N', async () => {
|
||||
mockGenerateContent
|
||||
.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-A' }] } }],
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-B' }] } }],
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-C' }] } }],
|
||||
})
|
||||
.mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-D' }] } }],
|
||||
});
|
||||
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
contextTurns: 2,
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['msg1']));
|
||||
await rewriter.rewrite(makeTurn(['msg2']));
|
||||
await rewriter.rewrite(makeTurn(['msg3']));
|
||||
|
||||
// 4th call — should include rewrite-B and rewrite-C (last 2), not rewrite-A
|
||||
await rewriter.rewrite(makeTurn(['msg4']));
|
||||
const input =
|
||||
mockGenerateContent.mock.calls[3][0].contents[0].parts[0].text;
|
||||
expect(input).not.toContain('rewrite-A');
|
||||
expect(input).toContain('rewrite-B');
|
||||
expect(input).toContain('rewrite-C');
|
||||
});
|
||||
|
||||
it('should include all rewrites when contextTurns="all"', async () => {
|
||||
mockGenerateContent
|
||||
.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-1' }] } }],
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-2' }] } }],
|
||||
})
|
||||
.mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'rewrite-3' }] } }],
|
||||
});
|
||||
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
contextTurns: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['msg1']));
|
||||
await rewriter.rewrite(makeTurn(['msg2']));
|
||||
await rewriter.rewrite(makeTurn(['msg3']));
|
||||
|
||||
const input =
|
||||
mockGenerateContent.mock.calls[2][0].contents[0].parts[0].text;
|
||||
expect(input).toContain('rewrite-1');
|
||||
expect(input).toContain('rewrite-2');
|
||||
});
|
||||
});
|
||||
|
||||
describe('model override', () => {
|
||||
it('should use rewriteConfig.model when set', async () => {
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
model: 'custom-rewrite-model',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['hello']));
|
||||
expect(mockGenerateContent.mock.calls[0][0].model).toBe(
|
||||
'custom-rewrite-model',
|
||||
);
|
||||
});
|
||||
|
||||
it('should fall back to config.getModel() when model is empty', async () => {
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['hello']));
|
||||
expect(mockGenerateContent.mock.calls[0][0].model).toBe('test-model');
|
||||
});
|
||||
});
|
||||
|
||||
describe('filtering', () => {
|
||||
it('should return null for empty input', async () => {
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
const result = await rewriter.rewrite(makeTurn([], []));
|
||||
expect(result).toBeNull();
|
||||
expect(mockGenerateContent).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return null when LLM returns short text', async () => {
|
||||
mockGenerateContent.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'hi' }] } }],
|
||||
});
|
||||
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
const result = await rewriter.rewrite(makeTurn(['some input text here']));
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should not accumulate failed rewrites in history', async () => {
|
||||
mockGenerateContent.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: '' }] } }],
|
||||
});
|
||||
mockGenerateContent.mockResolvedValueOnce({
|
||||
candidates: [{ content: { parts: [{ text: 'second rewrite ok' }] } }],
|
||||
});
|
||||
|
||||
const rewriter = new LlmRewriter(makeConfig(), {
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
} as MessageRewriteConfig);
|
||||
|
||||
await rewriter.rewrite(makeTurn(['first'])); // returns null
|
||||
await rewriter.rewrite(makeTurn(['second']));
|
||||
|
||||
// Second call should have no context (first rewrite returned null)
|
||||
const input =
|
||||
mockGenerateContent.mock.calls[1][0].contents[0].parts[0].text;
|
||||
expect(input).not.toContain('上一轮改写结果');
|
||||
});
|
||||
});
|
||||
});
|
||||
176
packages/cli/src/acp-integration/session/rewrite/LlmRewriter.ts
Normal file
176
packages/cli/src/acp-integration/session/rewrite/LlmRewriter.ts
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { readFileSync, existsSync } from 'node:fs';
|
||||
import { resolve } from 'node:path';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { createDebugLogger } from '@qwen-code/qwen-code-core';
|
||||
import type { TurnContent, MessageRewriteConfig } from './types.js';
|
||||
|
||||
const debugLogger = createDebugLogger('MESSAGE_REWRITER');
|
||||
|
||||
const DEFAULT_REWRITE_PROMPT = `You are an assistant that rewrites raw coding-agent output into concise, user-friendly progress updates.
|
||||
|
||||
The agent is a software engineering assistant that reads files, writes code, runs commands, and uses tools. Its raw output mixes internal reasoning with user-facing information. Your job: extract what the user cares about, drop what they don't.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Strictly based on original**: only surface information already in the input. Never invent details, plans, or conclusions the agent didn't state.
|
||||
2. **Keep**: goals, decisions, key findings, results, errors that affect the user, status updates.
|
||||
3. **Drop**: file paths, tool/skill names, internal reasoning about which tool to call, code snippets, stack traces, "let me…" / "now I'll…" filler phrases.
|
||||
4. **Progress turns**: if the agent is just starting a step (reading files, running a command, exploring code), output one short sentence describing what's happening — so the user isn't staring at silence.
|
||||
5. **Rich content**: if the input already contains well-structured user-facing content (tables, lists, formatted results), do light cleanup only (remove stray paths/tool names) and preserve the structure.
|
||||
6. **Pure internal ops**: if the input is entirely internal (fixing a typo in its own code, retrying a failed tool call, creating temp directories) → return empty string.
|
||||
7. **Preserve data exactly**: never alter numbers, percentages, file sizes, error codes, or quoted output.
|
||||
|
||||
## Context continuity
|
||||
|
||||
If "Previous rewrite output" is provided, the user has already seen it. Don't repeat — build on it. If this turn adds nothing new, return empty string.
|
||||
|
||||
Output only the rewritten text, or empty string if the input has no user-facing value.`;
|
||||
|
||||
/**
|
||||
* Uses LLM to rewrite turn content into business-friendly text.
|
||||
* Called at the end of each model turn (after all chunks accumulated).
|
||||
*/
|
||||
export class LlmRewriter {
|
||||
private readonly prompt: string;
|
||||
/** Previous successful rewrite outputs, used as context for coherence */
|
||||
private outputHistory: string[] = [];
|
||||
/** How many previous outputs to include: 0=none, N=last N, Infinity=all */
|
||||
private readonly contextTurns: number;
|
||||
|
||||
private readonly rewriteModel: string | undefined;
|
||||
|
||||
constructor(
|
||||
private readonly config: Config,
|
||||
rewriteConfig: MessageRewriteConfig,
|
||||
) {
|
||||
this.rewriteModel = rewriteConfig.model || undefined;
|
||||
this.contextTurns =
|
||||
rewriteConfig.contextTurns === 'all'
|
||||
? Infinity
|
||||
: (rewriteConfig.contextTurns ?? 1);
|
||||
// promptFile takes precedence over inline prompt
|
||||
if (rewriteConfig.promptFile) {
|
||||
const filePath = resolve(rewriteConfig.promptFile);
|
||||
if (existsSync(filePath)) {
|
||||
this.prompt = readFileSync(filePath, 'utf-8').trim();
|
||||
debugLogger.info(
|
||||
`Loaded rewrite prompt from file: ${filePath} (${this.prompt.length} chars)`,
|
||||
);
|
||||
} else {
|
||||
debugLogger.warn(
|
||||
`Rewrite prompt file not found: ${filePath}, using default`,
|
||||
);
|
||||
this.prompt = DEFAULT_REWRITE_PROMPT;
|
||||
}
|
||||
} else {
|
||||
this.prompt = rewriteConfig.prompt || DEFAULT_REWRITE_PROMPT;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Rewrite a turn's content using LLM.
|
||||
* Returns null if the turn has no valuable content for users.
|
||||
*/
|
||||
async rewrite(
|
||||
turnContent: TurnContent,
|
||||
signal?: AbortSignal,
|
||||
): Promise<string | null> {
|
||||
// Build input text from turn content
|
||||
const inputParts: string[] = [];
|
||||
|
||||
if (turnContent.thoughts.length > 0) {
|
||||
inputParts.push('[内部推理]\n' + turnContent.thoughts.join('\n'));
|
||||
}
|
||||
if (turnContent.messages.length > 0) {
|
||||
inputParts.push('[回复文本]\n' + turnContent.messages.join('\n'));
|
||||
}
|
||||
|
||||
// Prepend previous rewrite outputs as context for coherence
|
||||
if (this.contextTurns > 0 && this.outputHistory.length > 0) {
|
||||
const contextSlice =
|
||||
this.contextTurns === Infinity
|
||||
? this.outputHistory
|
||||
: this.outputHistory.slice(-this.contextTurns);
|
||||
inputParts.unshift('[上一轮改写结果]\n' + contextSlice.join('\n---\n'));
|
||||
}
|
||||
|
||||
const inputText = inputParts.join('\n\n');
|
||||
if (!inputText.trim()) return null;
|
||||
|
||||
// Skip very short turns that are likely just transitions
|
||||
if (inputText.length < 10) return null;
|
||||
|
||||
debugLogger.info(
|
||||
`[REWRITE INPUT] system_prompt_len=${this.prompt.length} input_len=${inputText.length} context_turns=${this.outputHistory.length}\n` +
|
||||
`--- INPUT TEXT ---\n${inputText}\n---`,
|
||||
);
|
||||
|
||||
try {
|
||||
const contentGenerator = this.config.getContentGenerator();
|
||||
if (!contentGenerator) {
|
||||
debugLogger.warn('No content generator available for rewriting');
|
||||
return null;
|
||||
}
|
||||
|
||||
const model = this.rewriteModel || this.config.getModel();
|
||||
|
||||
const result = await contentGenerator.generateContent(
|
||||
{
|
||||
model,
|
||||
config: {
|
||||
systemInstruction: this.prompt,
|
||||
abortSignal: signal,
|
||||
temperature: 0.3,
|
||||
maxOutputTokens: 1024,
|
||||
// Disable thinking to avoid thinking leaking into output
|
||||
thinkingConfig: { includeThoughts: false },
|
||||
},
|
||||
contents: [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [{ text: inputText }],
|
||||
},
|
||||
],
|
||||
},
|
||||
`rewrite-turn`,
|
||||
);
|
||||
|
||||
// Extract only non-thought text parts
|
||||
const rewritten =
|
||||
result.candidates?.[0]?.content?.parts
|
||||
?.filter((p) => !p.thought)
|
||||
.map((p) => p.text)
|
||||
.filter(Boolean)
|
||||
.join('') ?? '';
|
||||
|
||||
// If LLM returns empty or very short, skip
|
||||
if (!rewritten.trim() || rewritten.trim().length < 5) {
|
||||
debugLogger.info(`[REWRITE OUTPUT] empty or too short, skipping`);
|
||||
return null;
|
||||
}
|
||||
|
||||
const trimmed = rewritten.trim();
|
||||
|
||||
debugLogger.info(
|
||||
`[REWRITE OUTPUT] len=${trimmed.length}\n` +
|
||||
`--- OUTPUT ---\n${trimmed}\n---`,
|
||||
);
|
||||
|
||||
// Update context for next turn
|
||||
this.outputHistory.push(trimmed);
|
||||
|
||||
return trimmed;
|
||||
} catch (error) {
|
||||
debugLogger.warn(
|
||||
`LLM rewrite failed, skipping: ${error instanceof Error ? error.message : String(error)}`,
|
||||
);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,206 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import type { SessionUpdate } from '@agentclientprotocol/sdk';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
|
||||
// Mock core to avoid Vite https resolution issue
|
||||
vi.mock('@qwen-code/qwen-code-core', () => ({
|
||||
createDebugLogger: () => ({
|
||||
info: vi.fn(),
|
||||
warn: vi.fn(),
|
||||
debug: vi.fn(),
|
||||
error: vi.fn(),
|
||||
}),
|
||||
}));
|
||||
|
||||
// Mock LlmRewriter to avoid real LLM calls
|
||||
vi.mock('./LlmRewriter.js', () => ({
|
||||
LlmRewriter: vi.fn().mockImplementation(() => ({
|
||||
rewrite: vi.fn().mockResolvedValue('rewritten text'),
|
||||
})),
|
||||
}));
|
||||
|
||||
// Import after mocks are set up
|
||||
const { MessageRewriteMiddleware } = await import(
|
||||
'./MessageRewriteMiddleware.js'
|
||||
);
|
||||
|
||||
function createMiddleware(
|
||||
target: 'message' | 'thought' | 'all' = 'all',
|
||||
sendUpdate?: ReturnType<typeof vi.fn>,
|
||||
) {
|
||||
const mockSendUpdate = sendUpdate ?? vi.fn().mockResolvedValue(undefined);
|
||||
const middleware = new MessageRewriteMiddleware(
|
||||
{} as Config,
|
||||
{ enabled: true, target, prompt: 'test prompt' },
|
||||
mockSendUpdate,
|
||||
);
|
||||
return { middleware, mockSendUpdate };
|
||||
}
|
||||
|
||||
describe('MessageRewriteMiddleware', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('interceptUpdate — pass-through', () => {
|
||||
it('should pass through non-message updates unchanged', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware();
|
||||
const update = {
|
||||
sessionUpdate: 'tool_call_update',
|
||||
content: { text: 'progress' },
|
||||
} as unknown as SessionUpdate;
|
||||
|
||||
await middleware.interceptUpdate(update);
|
||||
expect(mockSendUpdate).toHaveBeenCalledWith(update);
|
||||
});
|
||||
|
||||
it('should always send original message/thought as-is', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware();
|
||||
const msgUpdate = {
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'hello' },
|
||||
} as unknown as SessionUpdate;
|
||||
|
||||
await middleware.interceptUpdate(msgUpdate);
|
||||
expect(mockSendUpdate).toHaveBeenCalledWith(msgUpdate);
|
||||
});
|
||||
});
|
||||
|
||||
describe('interceptUpdate — target filtering', () => {
|
||||
it('should accumulate messages when target is "message"', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware('message');
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'msg' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_thought_chunk',
|
||||
content: { type: 'text', text: 'thought' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
// Flush and wait
|
||||
await middleware.flushTurn();
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
// Original pass-through (2) + rewritten (1)
|
||||
expect(mockSendUpdate).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('should not accumulate thoughts when target is "message"', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware('message');
|
||||
|
||||
// Only thought, no message — flush should produce nothing
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_thought_chunk',
|
||||
content: { type: 'text', text: 'thought only' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.flushTurn();
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
// Only the original pass-through, no rewrite
|
||||
expect(mockSendUpdate).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should accumulate both when target is "both"', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware('all');
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'msg' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_thought_chunk',
|
||||
content: { type: 'text', text: 'thought' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.flushTurn();
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
// 2 pass-throughs + 1 rewrite
|
||||
expect(mockSendUpdate).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe('flushTurn — tool_call boundary', () => {
|
||||
it('should flush before passing through tool_call', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware();
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'before tool' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'tool_call',
|
||||
callId: '123',
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
// pass-through msg + tool_call + rewrite
|
||||
expect(mockSendUpdate).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe('waitForPendingRewrites', () => {
|
||||
it('should wait for multiple pending rewrites', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware();
|
||||
|
||||
// Simulate 3 turns
|
||||
for (let i = 0; i < 3; i++) {
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: `turn ${i}` },
|
||||
} as unknown as SessionUpdate);
|
||||
await middleware.flushTurn();
|
||||
}
|
||||
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
// 3 pass-throughs + 3 rewrites
|
||||
expect(mockSendUpdate).toHaveBeenCalledTimes(6);
|
||||
});
|
||||
|
||||
it('should be safe to call when no rewrites are pending', async () => {
|
||||
const { middleware } = createMiddleware();
|
||||
await expect(
|
||||
middleware.waitForPendingRewrites(),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('rewrite metadata', () => {
|
||||
it('should emit rewritten message with _meta.rewritten=true', async () => {
|
||||
const { middleware, mockSendUpdate } = createMiddleware();
|
||||
|
||||
await middleware.interceptUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'content' },
|
||||
} as unknown as SessionUpdate);
|
||||
|
||||
await middleware.flushTurn();
|
||||
await middleware.waitForPendingRewrites();
|
||||
|
||||
const rewriteCall = mockSendUpdate.mock.calls.find(
|
||||
(call: unknown[]) =>
|
||||
(call[0] as Record<string, unknown>)['_meta'] !== undefined,
|
||||
);
|
||||
expect(rewriteCall).toBeDefined();
|
||||
const meta = (rewriteCall![0] as Record<string, unknown>)[
|
||||
'_meta'
|
||||
] as Record<string, unknown>;
|
||||
expect(meta['rewritten']).toBe(true);
|
||||
expect(meta['turnIndex']).toBe(1);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,159 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SessionUpdate } from '@agentclientprotocol/sdk';
|
||||
import type { Config } from '@qwen-code/qwen-code-core';
|
||||
import { createDebugLogger } from '@qwen-code/qwen-code-core';
|
||||
import type { MessageRewriteConfig } from './types.js';
|
||||
import { TurnBuffer } from './TurnBuffer.js';
|
||||
import { LlmRewriter } from './LlmRewriter.js';
|
||||
|
||||
const debugLogger = createDebugLogger('MESSAGE_REWRITE');
|
||||
|
||||
/**
|
||||
* Middleware that intercepts ACP messages and appends LLM-rewritten
|
||||
* versions with _meta.rewritten=true.
|
||||
*
|
||||
* Original messages are sent as-is (no modification).
|
||||
* At the end of each turn, a rewritten message is appended.
|
||||
*
|
||||
* Flow:
|
||||
* 1. Original chunks pass through unmodified
|
||||
* 2. Chunks are accumulated in TurnBuffer
|
||||
* 3. When a turn ends (tool_call starts, or session ends),
|
||||
* LlmRewriter rewrites the accumulated content
|
||||
* 4. Rewritten text is emitted as agent_message_chunk with _meta.rewritten=true
|
||||
*/
|
||||
export class MessageRewriteMiddleware {
|
||||
private readonly turnBuffer: TurnBuffer;
|
||||
private readonly rewriter: LlmRewriter;
|
||||
private readonly target: MessageRewriteConfig['target'];
|
||||
private turnIndex = 0;
|
||||
|
||||
constructor(
|
||||
config: Config,
|
||||
rewriteConfig: MessageRewriteConfig,
|
||||
private readonly sendUpdate: (update: SessionUpdate) => Promise<void>,
|
||||
) {
|
||||
this.turnBuffer = new TurnBuffer();
|
||||
this.rewriter = new LlmRewriter(config, rewriteConfig);
|
||||
this.target = rewriteConfig.target;
|
||||
}
|
||||
|
||||
/**
|
||||
* Intercept an ACP update. Original messages pass through,
|
||||
* thought/message chunks are also accumulated for turn-end rewriting.
|
||||
*/
|
||||
async interceptUpdate(
|
||||
update: SessionUpdate,
|
||||
signal?: AbortSignal,
|
||||
): Promise<void> {
|
||||
const updateRecord = update as Record<string, unknown>;
|
||||
const updateType = updateRecord['sessionUpdate'] as string;
|
||||
|
||||
// tool_call signals turn boundary — flush before passing through
|
||||
if (updateType === 'tool_call') {
|
||||
await this.flushTurn(signal);
|
||||
this.turnBuffer.markToolCall();
|
||||
return this.sendUpdate(update);
|
||||
}
|
||||
|
||||
// tool_call_update, plan, available_commands, etc. → pass through
|
||||
if (
|
||||
updateType !== 'agent_thought_chunk' &&
|
||||
updateType !== 'agent_message_chunk'
|
||||
) {
|
||||
return this.sendUpdate(update);
|
||||
}
|
||||
|
||||
const content = updateRecord['content'] as
|
||||
| Record<string, string>
|
||||
| undefined;
|
||||
const text = content?.['text'] ?? '';
|
||||
|
||||
// Always send original message as-is
|
||||
await this.sendUpdate(update);
|
||||
|
||||
// Accumulate for turn-end rewriting
|
||||
if (updateType === 'agent_thought_chunk') {
|
||||
if (this.target === 'thought' || this.target === 'all') {
|
||||
this.turnBuffer.appendThought(text);
|
||||
}
|
||||
} else if (updateType === 'agent_message_chunk') {
|
||||
if (this.target === 'message' || this.target === 'all') {
|
||||
this.turnBuffer.appendMessage(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Pending rewrite promises — all must settle before session exits */
|
||||
private pendingRewrites: Array<Promise<void>> = [];
|
||||
|
||||
/**
|
||||
* Flush the turn buffer: rewrite accumulated content and emit.
|
||||
*
|
||||
* Non-blocking: rewrite runs in background, parallel to tool execution.
|
||||
*
|
||||
* Called when:
|
||||
* - A tool_call is about to be emitted (turn boundary)
|
||||
* - Usage metadata is emitted (end of model response)
|
||||
* - Session prompt ends
|
||||
*/
|
||||
async flushTurn(signal?: AbortSignal): Promise<void> {
|
||||
const content = this.turnBuffer.flush();
|
||||
if (!content) return;
|
||||
|
||||
this.turnIndex++;
|
||||
const turnIdx = this.turnIndex;
|
||||
|
||||
// Always enforce a 30s timeout, combined with caller's signal if provided
|
||||
const timeoutSignal = AbortSignal.timeout(30_000);
|
||||
const rewriteSignal = signal
|
||||
? AbortSignal.any([signal, timeoutSignal])
|
||||
: timeoutSignal;
|
||||
|
||||
this.pendingRewrites.push(
|
||||
(async () => {
|
||||
try {
|
||||
const rewritten = await this.rewriter.rewrite(content, rewriteSignal);
|
||||
if (!rewritten) {
|
||||
debugLogger.info(`Turn ${turnIdx}: no rewrite output`);
|
||||
return;
|
||||
}
|
||||
|
||||
debugLogger.info(
|
||||
`Turn ${turnIdx}: rewritten ${rewritten.length} chars`,
|
||||
);
|
||||
|
||||
// Emit rewritten message with special _meta
|
||||
await this.sendUpdate({
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: rewritten },
|
||||
_meta: {
|
||||
rewritten: true,
|
||||
turnIndex: turnIdx,
|
||||
},
|
||||
} as SessionUpdate);
|
||||
} catch (error) {
|
||||
debugLogger.warn(
|
||||
`Turn ${turnIdx}: rewrite failed: ${error instanceof Error ? error.message : String(error)}`,
|
||||
);
|
||||
}
|
||||
})(),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait for all pending rewrites to complete.
|
||||
* Call this before session ends to ensure all rewrites are flushed.
|
||||
*/
|
||||
async waitForPendingRewrites(): Promise<void> {
|
||||
if (this.pendingRewrites.length > 0) {
|
||||
await Promise.allSettled(this.pendingRewrites);
|
||||
this.pendingRewrites = [];
|
||||
}
|
||||
}
|
||||
}
|
||||
32
packages/cli/src/acp-integration/session/rewrite/README.md
Normal file
32
packages/cli/src/acp-integration/session/rewrite/README.md
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
# Message Rewrite Middleware
|
||||
|
||||
> **⚠️ Temporary Solution — subject to change or removal at any time.**
|
||||
>
|
||||
> This is a stopgap implementation. We are considering a hook-based approach that would be more decoupled and extensible. Ideas and suggestions for a better design are very welcome.
|
||||
|
||||
## Use Case
|
||||
|
||||
When a coding agent is integrated into vertical business scenarios (data analysis, ops, report generation, etc.), the raw output often contains technical details (file paths, tool calls, internal reasoning) that end users don't care about. By configuring a rewrite prompt, the output can be transformed into business-friendly language.
|
||||
|
||||
## How It Works
|
||||
|
||||
1. Original messages are **passed through as-is** — no modification
|
||||
2. At the end of each turn (before tool calls / at response end), accumulated thought + message chunks are sent to a separate LLM call for rewriting
|
||||
3. Rewritten text is appended as a new `agent_message_chunk` with `_meta.rewritten: true`
|
||||
4. The client decides which version to display based on `_meta.rewritten`
|
||||
|
||||
## Configuration
|
||||
|
||||
Add to `settings.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"messageRewrite": {
|
||||
"enabled": true,
|
||||
"target": "all",
|
||||
"promptFile": ".qwen/rewrite-prompt.txt",
|
||||
"model": "qwen3-plus",
|
||||
"contextTurns": 1
|
||||
}
|
||||
}
|
||||
```
|
||||
|
|
@ -0,0 +1,112 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { TurnBuffer } from './TurnBuffer.js';
|
||||
|
||||
describe('TurnBuffer', () => {
|
||||
let buffer: TurnBuffer;
|
||||
|
||||
beforeEach(() => {
|
||||
buffer = new TurnBuffer();
|
||||
});
|
||||
|
||||
describe('isEmpty', () => {
|
||||
it('should be empty initially', () => {
|
||||
expect(buffer.isEmpty).toBe(true);
|
||||
});
|
||||
|
||||
it('should not be empty after appending a message', () => {
|
||||
buffer.appendMessage('hello');
|
||||
expect(buffer.isEmpty).toBe(false);
|
||||
});
|
||||
|
||||
it('should not be empty after appending a thought', () => {
|
||||
buffer.appendThought('thinking...');
|
||||
expect(buffer.isEmpty).toBe(false);
|
||||
});
|
||||
|
||||
it('should be empty after flush', () => {
|
||||
buffer.appendMessage('hello');
|
||||
buffer.flush();
|
||||
expect(buffer.isEmpty).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('appendMessage / appendThought', () => {
|
||||
it('should ignore empty strings', () => {
|
||||
buffer.appendMessage('');
|
||||
buffer.appendThought('');
|
||||
expect(buffer.isEmpty).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('markToolCall', () => {
|
||||
it('should set hasToolCalls in flushed content', () => {
|
||||
buffer.appendMessage('text');
|
||||
buffer.markToolCall();
|
||||
const content = buffer.flush();
|
||||
expect(content?.hasToolCalls).toBe(true);
|
||||
});
|
||||
|
||||
it('should default hasToolCalls to false', () => {
|
||||
buffer.appendMessage('text');
|
||||
const content = buffer.flush();
|
||||
expect(content?.hasToolCalls).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('flush', () => {
|
||||
it('should return null when buffer is empty', () => {
|
||||
expect(buffer.flush()).toBeNull();
|
||||
});
|
||||
|
||||
it('should return null when only whitespace was appended', () => {
|
||||
buffer.appendMessage(' ');
|
||||
buffer.appendThought(' \n ');
|
||||
expect(buffer.flush()).toBeNull();
|
||||
});
|
||||
|
||||
it('should return accumulated messages and thoughts', () => {
|
||||
buffer.appendThought('thought 1');
|
||||
buffer.appendThought('thought 2');
|
||||
buffer.appendMessage('msg 1');
|
||||
buffer.appendMessage('msg 2');
|
||||
|
||||
const content = buffer.flush();
|
||||
expect(content).toEqual({
|
||||
thoughts: ['thought 1', 'thought 2'],
|
||||
messages: ['msg 1', 'msg 2'],
|
||||
hasToolCalls: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should filter out whitespace-only entries', () => {
|
||||
buffer.appendThought(' ');
|
||||
buffer.appendThought('real thought');
|
||||
buffer.appendMessage('');
|
||||
buffer.appendMessage('real message');
|
||||
|
||||
const content = buffer.flush();
|
||||
expect(content?.thoughts).toEqual(['real thought']);
|
||||
expect(content?.messages).toEqual(['real message']);
|
||||
});
|
||||
|
||||
it('should reset buffer after flush', () => {
|
||||
buffer.appendMessage('first');
|
||||
buffer.markToolCall();
|
||||
buffer.flush();
|
||||
|
||||
buffer.appendMessage('second');
|
||||
const content = buffer.flush();
|
||||
expect(content).toEqual({
|
||||
thoughts: [],
|
||||
messages: ['second'],
|
||||
hasToolCalls: false,
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { TurnContent } from './types.js';
|
||||
|
||||
/**
|
||||
* Accumulates thought and message chunks for a single model turn.
|
||||
* A turn ends when tool calls begin or the model stops generating.
|
||||
*/
|
||||
export class TurnBuffer {
|
||||
private thoughts: string[] = [];
|
||||
private messages: string[] = [];
|
||||
private _hasToolCalls = false;
|
||||
|
||||
appendThought(text: string): void {
|
||||
if (text) this.thoughts.push(text);
|
||||
}
|
||||
|
||||
appendMessage(text: string): void {
|
||||
if (text) this.messages.push(text);
|
||||
}
|
||||
|
||||
markToolCall(): void {
|
||||
this._hasToolCalls = true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns accumulated content and resets the buffer.
|
||||
* Returns null if buffer is empty.
|
||||
*/
|
||||
flush(): TurnContent | null {
|
||||
const thoughtText = this.thoughts.join('');
|
||||
const messageText = this.messages.join('');
|
||||
|
||||
if (!thoughtText.trim() && !messageText.trim()) {
|
||||
this.reset();
|
||||
return null;
|
||||
}
|
||||
|
||||
const content: TurnContent = {
|
||||
thoughts: this.thoughts.filter((t) => t.trim()),
|
||||
messages: this.messages.filter((m) => m.trim()),
|
||||
hasToolCalls: this._hasToolCalls,
|
||||
};
|
||||
|
||||
this.reset();
|
||||
return content;
|
||||
}
|
||||
|
||||
private reset(): void {
|
||||
this.thoughts = [];
|
||||
this.messages = [];
|
||||
this._hasToolCalls = false;
|
||||
}
|
||||
|
||||
get isEmpty(): boolean {
|
||||
return this.thoughts.length === 0 && this.messages.length === 0;
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { loadRewriteConfig } from './config.js';
|
||||
import type { LoadedSettings } from '../../../config/settings.js';
|
||||
|
||||
/**
|
||||
* Build a minimal LoadedSettings stub with only the fields
|
||||
* that loadRewriteConfig actually reads (user/workspace originalSettings + isTrusted).
|
||||
*/
|
||||
function makeSettings(
|
||||
overrides: {
|
||||
userRewrite?: Record<string, unknown>;
|
||||
workspaceRewrite?: Record<string, unknown>;
|
||||
isTrusted?: boolean;
|
||||
} = {},
|
||||
): LoadedSettings {
|
||||
return {
|
||||
user: {
|
||||
originalSettings: overrides.userRewrite
|
||||
? { messageRewrite: overrides.userRewrite }
|
||||
: {},
|
||||
},
|
||||
workspace: {
|
||||
originalSettings: overrides.workspaceRewrite
|
||||
? { messageRewrite: overrides.workspaceRewrite }
|
||||
: {},
|
||||
},
|
||||
isTrusted: overrides.isTrusted ?? true,
|
||||
} as unknown as LoadedSettings;
|
||||
}
|
||||
|
||||
describe('loadRewriteConfig', () => {
|
||||
it('should return undefined when no config is set', () => {
|
||||
const settings = makeSettings();
|
||||
expect(loadRewriteConfig(settings)).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return user config when only user config is set', () => {
|
||||
const settings = makeSettings({
|
||||
userRewrite: { enabled: true, target: 'all', prompt: 'user prompt' },
|
||||
});
|
||||
const config = loadRewriteConfig(settings);
|
||||
expect(config).toEqual({
|
||||
enabled: true,
|
||||
target: 'all',
|
||||
prompt: 'user prompt',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return workspace config when trusted', () => {
|
||||
const settings = makeSettings({
|
||||
userRewrite: { enabled: false, target: 'message' },
|
||||
workspaceRewrite: { enabled: true, target: 'all', prompt: 'ws prompt' },
|
||||
isTrusted: true,
|
||||
});
|
||||
const config = loadRewriteConfig(settings);
|
||||
expect(config?.enabled).toBe(true);
|
||||
expect(config?.prompt).toBe('ws prompt');
|
||||
});
|
||||
|
||||
it('should ignore workspace config when untrusted', () => {
|
||||
const settings = makeSettings({
|
||||
userRewrite: { enabled: false, target: 'message' },
|
||||
workspaceRewrite: { enabled: true, target: 'all', prompt: 'malicious' },
|
||||
isTrusted: false,
|
||||
});
|
||||
const config = loadRewriteConfig(settings);
|
||||
expect(config?.enabled).toBe(false);
|
||||
expect(config?.prompt).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should fall back to user config when workspace has no rewrite config', () => {
|
||||
const settings = makeSettings({
|
||||
userRewrite: { enabled: true, target: 'thought' },
|
||||
isTrusted: true,
|
||||
});
|
||||
const config = loadRewriteConfig(settings);
|
||||
expect(config?.target).toBe('thought');
|
||||
});
|
||||
});
|
||||
28
packages/cli/src/acp-integration/session/rewrite/config.ts
Normal file
28
packages/cli/src/acp-integration/session/rewrite/config.ts
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { LoadedSettings } from '../../../config/settings.js';
|
||||
import type { MessageRewriteConfig } from './types.js';
|
||||
|
||||
/**
|
||||
* Reads messageRewrite configuration from user/workspace originalSettings.
|
||||
* Workspace settings are only used when the workspace is trusted,
|
||||
* preventing untrusted repos from enabling the rewriter with a custom prompt.
|
||||
*/
|
||||
export function loadRewriteConfig(
|
||||
settings: LoadedSettings,
|
||||
): MessageRewriteConfig | undefined {
|
||||
const userOriginal = settings.user?.originalSettings as
|
||||
| Record<string, unknown>
|
||||
| undefined;
|
||||
const workspaceOriginal = settings.isTrusted
|
||||
? (settings.workspace?.originalSettings as
|
||||
| Record<string, unknown>
|
||||
| undefined)
|
||||
: undefined;
|
||||
return (workspaceOriginal?.['messageRewrite'] ??
|
||||
userOriginal?.['messageRewrite']) as MessageRewriteConfig | undefined;
|
||||
}
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
export { MessageRewriteMiddleware } from './MessageRewriteMiddleware.js';
|
||||
export { loadRewriteConfig } from './config.js';
|
||||
export type { MessageRewriteConfig, TurnContent } from './types.js';
|
||||
36
packages/cli/src/acp-integration/session/rewrite/types.ts
Normal file
36
packages/cli/src/acp-integration/session/rewrite/types.ts
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Qwen
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* Configuration for ACP message rewriting.
|
||||
* Loaded from .qwen/settings.json under "messageRewrite" key.
|
||||
*/
|
||||
export interface MessageRewriteConfig {
|
||||
/** Whether message rewriting is enabled */
|
||||
enabled: boolean;
|
||||
/** Which message types to rewrite */
|
||||
target: 'message' | 'thought' | 'all';
|
||||
/** LLM rewrite prompt (system prompt for the rewriter). Inline string. */
|
||||
prompt?: string;
|
||||
/** Path to a file containing the rewrite prompt. Resolved relative to CWD.
|
||||
* Takes precedence over `prompt` if both are set. */
|
||||
promptFile?: string;
|
||||
/** Model to use for rewriting (empty = use current model) */
|
||||
model?: string;
|
||||
/** Number of previous rewrite outputs to include as context.
|
||||
* 1 = last rewrite only (default), "all" = all previous rewrites,
|
||||
* 0 = no context, N = last N rewrites. */
|
||||
contextTurns?: number | 'all';
|
||||
}
|
||||
|
||||
/**
|
||||
* Accumulated content for a single turn.
|
||||
*/
|
||||
export interface TurnContent {
|
||||
thoughts: string[];
|
||||
messages: string[];
|
||||
hasToolCalls: boolean;
|
||||
}
|
||||
|
|
@ -11,6 +11,7 @@ import type {
|
|||
ToolCallLocation,
|
||||
ToolKind,
|
||||
} from '@agentclientprotocol/sdk';
|
||||
import type { MessageRewriteMiddleware } from './rewrite/index.js';
|
||||
|
||||
export type ApprovalModeValue = 'plan' | 'default' | 'auto-edit' | 'yolo';
|
||||
|
||||
|
|
@ -29,6 +30,9 @@ export interface SessionUpdateSender {
|
|||
export interface SessionContext extends SessionUpdateSender {
|
||||
readonly sessionId: string;
|
||||
readonly config: Config;
|
||||
/** Optional message rewrite middleware for ACP message transformation.
|
||||
* Installed after history replay to avoid rewriting historical messages. */
|
||||
messageRewriter?: MessageRewriteMiddleware;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue