Skip to content

Commit 0a02ec3

Browse files
jonathanpeterwujonathanpeterwu
committed
feat(sweep): add Sweep Next-Edit wrapper with llama-server support
Phase A implementation of Sweep code completion predictions: - Server manager for llama-server lifecycle (start/stop/status) - Prompt builder in Sweep <|file_sep|> format - HTTP prediction client for OpenAI-compatible API - CLI commands: stackmemory sweep start/stop/status/predict/hook Integrates with existing post-edit-sweep.js hook for Claude Code.
1 parent 9343a53 commit 0a02ec3

7 files changed

Lines changed: 1035 additions & 0 deletions

File tree

‎src/cli/commands/sweep.ts‎

Lines changed: 297 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,297 @@
1+
/**
2+
* Sweep CLI Command
3+
* Manage Sweep Next-Edit prediction server
4+
*/
5+
6+
import { Command } from 'commander';
7+
import chalk from 'chalk';
8+
import ora from 'ora';
9+
import { existsSync } from 'fs';
10+
import { join } from 'path';
11+
import {
12+
createServerManager,
13+
createPredictionClient,
14+
SweepServerConfig,
15+
DEFAULT_SERVER_CONFIG,
16+
} from '../../features/sweep/index.js';
17+
18+
const HOME = process.env['HOME'] || '/tmp';
19+
20+
export function createSweepCommand(): Command {
21+
const cmd = new Command('sweep')
22+
.description('Manage Sweep Next-Edit prediction server')
23+
.addHelpText(
24+
'after',
25+
`
26+
Examples:
27+
stackmemory sweep start Start the Sweep server
28+
stackmemory sweep stop Stop the Sweep server
29+
stackmemory sweep status Check server status
30+
stackmemory sweep predict Run a prediction manually
31+
stackmemory sweep hook Install/check Claude Code hook
32+
33+
The Sweep server uses the Sweep 1.5B model for next-edit predictions.
34+
Predictions are triggered after file edits via Claude Code hooks.
35+
`
36+
);
37+
38+
// Start command
39+
cmd
40+
.command('start')
41+
.description('Start the Sweep prediction server')
42+
.option(
43+
'--port <number>',
44+
'Server port',
45+
String(DEFAULT_SERVER_CONFIG.port)
46+
)
47+
.option('--model <path>', 'Path to GGUF model file')
48+
.option(
49+
'--context <size>',
50+
'Context size',
51+
String(DEFAULT_SERVER_CONFIG.contextSize)
52+
)
53+
.option('--gpu-layers <n>', 'Number of GPU layers (0 for CPU only)', '0')
54+
.action(async (options) => {
55+
const spinner = ora('Starting Sweep server...').start();
56+
57+
try {
58+
const config: Partial<SweepServerConfig> = {
59+
port: parseInt(options.port, 10),
60+
contextSize: parseInt(options.context, 10),
61+
gpuLayers: parseInt(options.gpuLayers, 10),
62+
};
63+
64+
if (options.model) {
65+
config.modelPath = options.model;
66+
}
67+
68+
const manager = createServerManager(config);
69+
const status = await manager.startServer();
70+
71+
spinner.succeed(chalk.green('Sweep server started'));
72+
console.log(chalk.gray(` PID: ${status.pid}`));
73+
console.log(chalk.gray(` Port: ${status.port}`));
74+
console.log(chalk.gray(` Model: ${status.modelPath}`));
75+
} catch (error) {
76+
spinner.fail(chalk.red('Failed to start Sweep server'));
77+
console.error(chalk.red((error as Error).message));
78+
process.exit(1);
79+
}
80+
});
81+
82+
// Stop command
83+
cmd
84+
.command('stop')
85+
.description('Stop the Sweep prediction server')
86+
.action(async () => {
87+
const spinner = ora('Stopping Sweep server...').start();
88+
89+
try {
90+
const manager = createServerManager();
91+
await manager.stopServer();
92+
spinner.succeed(chalk.green('Sweep server stopped'));
93+
} catch (error) {
94+
spinner.fail(chalk.red('Failed to stop Sweep server'));
95+
console.error(chalk.red((error as Error).message));
96+
process.exit(1);
97+
}
98+
});
99+
100+
// Status command
101+
cmd
102+
.command('status')
103+
.description('Check Sweep server status')
104+
.action(async () => {
105+
try {
106+
const manager = createServerManager();
107+
const status = await manager.getStatus();
108+
109+
if (status.running) {
110+
console.log(chalk.green('Sweep server is running'));
111+
console.log(chalk.gray(` PID: ${status.pid}`));
112+
console.log(chalk.gray(` Port: ${status.port}`));
113+
console.log(chalk.gray(` Host: ${status.host}`));
114+
if (status.startedAt) {
115+
const uptime = Math.floor((Date.now() - status.startedAt) / 1000);
116+
console.log(chalk.gray(` Uptime: ${formatUptime(uptime)}`));
117+
}
118+
if (status.modelPath) {
119+
console.log(chalk.gray(` Model: ${status.modelPath}`));
120+
}
121+
} else {
122+
console.log(chalk.yellow('Sweep server is not running'));
123+
console.log(chalk.gray(' Start with: stackmemory sweep start'));
124+
}
125+
126+
// Check model availability
127+
const defaultModelPath = join(
128+
HOME,
129+
'.stackmemory',
130+
'models',
131+
'sweep',
132+
'sweep-next-edit-1.5b.q8_0.v2.gguf'
133+
);
134+
if (!existsSync(defaultModelPath)) {
135+
console.log('');
136+
console.log(chalk.yellow('Model not found at default location'));
137+
console.log(
138+
chalk.gray(
139+
' Download with:\n' +
140+
' huggingface-cli download sweepai/sweep-next-edit-1.5B \\\n' +
141+
' sweep-next-edit-1.5b.q8_0.v2.gguf \\\n' +
142+
' --local-dir ~/.stackmemory/models/sweep'
143+
)
144+
);
145+
}
146+
} catch (error) {
147+
console.error(
148+
chalk.red('Error checking status:'),
149+
(error as Error).message
150+
);
151+
process.exit(1);
152+
}
153+
});
154+
155+
// Predict command (for testing)
156+
cmd
157+
.command('predict')
158+
.description('Run a prediction manually (for testing)')
159+
.argument('<file>', 'File to predict edits for')
160+
.option(
161+
'--port <number>',
162+
'Server port',
163+
String(DEFAULT_SERVER_CONFIG.port)
164+
)
165+
.action(async (file, options) => {
166+
const spinner = ora('Running prediction...').start();
167+
168+
try {
169+
const { readFileSync } = await import('fs');
170+
const content = readFileSync(file, 'utf-8');
171+
172+
const client = createPredictionClient({
173+
port: parseInt(options.port, 10),
174+
});
175+
176+
// Check server
177+
const healthy = await client.checkHealth();
178+
if (!healthy) {
179+
spinner.fail(chalk.red('Server not running'));
180+
console.log(chalk.gray('Start with: stackmemory sweep start'));
181+
process.exit(1);
182+
}
183+
184+
const result = await client.predict({
185+
file_path: file,
186+
current_content: content,
187+
recent_diffs: [],
188+
});
189+
190+
spinner.stop();
191+
192+
if (result.success && result.predicted_content) {
193+
console.log(chalk.green('Prediction complete'));
194+
console.log(chalk.gray(`Latency: ${result.latency_ms}ms`));
195+
console.log(chalk.gray(`Tokens: ${result.tokens_generated}`));
196+
console.log('');
197+
console.log(chalk.cyan('Predicted content:'));
198+
console.log(result.predicted_content.slice(0, 500));
199+
if (result.predicted_content.length > 500) {
200+
console.log(chalk.gray('... (truncated)'));
201+
}
202+
} else if (result.success) {
203+
console.log(chalk.yellow('No changes predicted'));
204+
} else {
205+
console.log(chalk.red('Prediction failed:'), result.message);
206+
}
207+
} catch (error) {
208+
spinner.fail(chalk.red('Prediction failed'));
209+
console.error(chalk.red((error as Error).message));
210+
process.exit(1);
211+
}
212+
});
213+
214+
// Hook command
215+
cmd
216+
.command('hook')
217+
.description('Check Claude Code hook status')
218+
.action(async () => {
219+
const hookPath = join(HOME, '.claude', 'hooks', 'post-edit-sweep.js');
220+
const templatePath = join(
221+
process.cwd(),
222+
'templates',
223+
'claude-hooks',
224+
'post-edit-sweep.js'
225+
);
226+
227+
console.log(chalk.cyan('Sweep Hook Status'));
228+
console.log('');
229+
230+
if (existsSync(hookPath)) {
231+
console.log(chalk.green('Hook installed:'), hookPath);
232+
} else {
233+
console.log(chalk.yellow('Hook not installed'));
234+
console.log(chalk.gray(` Copy from: ${templatePath}`));
235+
console.log(chalk.gray(` To: ${hookPath}`));
236+
}
237+
238+
// Check settings.json
239+
const settingsPath = join(HOME, '.claude', 'settings.json');
240+
if (existsSync(settingsPath)) {
241+
try {
242+
const { readFileSync } = await import('fs');
243+
const settings = JSON.parse(readFileSync(settingsPath, 'utf-8'));
244+
const hasHook = JSON.stringify(settings).includes('post-edit-sweep');
245+
if (hasHook) {
246+
console.log(chalk.green('Hook registered in settings.json'));
247+
} else {
248+
console.log(chalk.yellow('Hook not registered in settings.json'));
249+
console.log(
250+
chalk.gray(
251+
' Add to hooks.PostToolUse:\n' +
252+
' { "matcher": "Edit", "hooks": [{ "type": "command", "command": "node ~/.claude/hooks/post-edit-sweep.js" }] }'
253+
)
254+
);
255+
}
256+
} catch {
257+
console.log(chalk.yellow('Could not read settings.json'));
258+
}
259+
} else {
260+
console.log(chalk.yellow('settings.json not found'));
261+
}
262+
263+
// Check state file
264+
const statePath = join(HOME, '.stackmemory', 'sweep-state.json');
265+
if (existsSync(statePath)) {
266+
try {
267+
const { readFileSync } = await import('fs');
268+
const state = JSON.parse(readFileSync(statePath, 'utf-8'));
269+
console.log(
270+
chalk.gray(
271+
` Recent diffs tracked: ${state.recentDiffs?.length || 0}`
272+
)
273+
);
274+
if (state.lastPrediction) {
275+
const ago = Math.floor(
276+
(Date.now() - state.lastPrediction.timestamp) / 1000
277+
);
278+
console.log(
279+
chalk.gray(` Last prediction: ${formatUptime(ago)} ago`)
280+
);
281+
}
282+
} catch {
283+
// Ignore
284+
}
285+
}
286+
});
287+
288+
return cmd;
289+
}
290+
291+
function formatUptime(seconds: number): string {
292+
if (seconds < 60) return `${seconds}s`;
293+
if (seconds < 3600) return `${Math.floor(seconds / 60)}m ${seconds % 60}s`;
294+
const hours = Math.floor(seconds / 3600);
295+
const mins = Math.floor((seconds % 3600) / 60);
296+
return `${hours}h ${mins}m`;
297+
}

‎src/cli/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import { registerSignupCommand } from './commands/signup.js';
4343
import { registerLogoutCommand, registerDbCommands } from './commands/db.js';
4444
import { createHooksCommand } from './commands/hooks.js';
4545
import { createDaemonCommand } from './commands/daemon.js';
46+
import { createSweepCommand } from './commands/sweep.js';
4647
import { createShellCommand } from './commands/shell.js';
4748
import { createAPICommand } from './commands/api.js';
4849
import { createCleanupProcessesCommand } from './commands/cleanup-processes.js';
@@ -712,6 +713,7 @@ if (isFeatureEnabled('ralph')) {
712713
});
713714
}
714715
program.addCommand(createDaemonCommand());
716+
program.addCommand(createSweepCommand());
715717
program.addCommand(createShellCommand());
716718
program.addCommand(createAPICommand());
717719
program.addCommand(createCleanupProcessesCommand());

‎src/features/sweep/index.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
/**
2+
* Sweep Next-Edit Feature
3+
*
4+
* Provides next-edit predictions using Sweep 1.5B model
5+
* via a local llama-server instance.
6+
*/
7+
8+
export * from './types.js';
9+
export * from './prompt-builder.js';
10+
export * from './prediction-client.js';
11+
export * from './sweep-server-manager.js';

0 commit comments

Comments
 (0)