feat: add Complete D-Bus method for code completion
Adds a Complete method that manages a dedicated copilot session, sends a FIM prompt with prefix/suffix context, waits for the result, and returns cleaned completion text. Includes unit tests for strip helpers.
This commit is contained in:
Levi Neely
parent
5afe7b69ac
commit
76c4ddcdd9
|
|
@ -0,0 +1,51 @@
|
|||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestStripCompletionNoise(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
in string
|
||||
out string
|
||||
}{
|
||||
{"plain", "hello\nworld", "hello\nworld"},
|
||||
{"markdown fences", "```go\nfmt.Println()\n```", "fmt.Println()"},
|
||||
{"leaked prefix tag", "<prefix>\ncode\n</prefix>", "code"},
|
||||
{"leaked suffix tag", "<suffix>\ncode\n</suffix>", "code"},
|
||||
{"info lines", ":: cost=0.01\nresult", "result"},
|
||||
{"mixed", "```\n:: info\nreal code\n<prefix>\n```", "real code"},
|
||||
{"empty", "", ""},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := stripCompletionNoise(tt.in)
|
||||
if got != tt.out {
|
||||
t.Errorf("stripCompletionNoise(%q) = %q, want %q", tt.in, got, tt.out)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripPrefixEcho(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prefix string
|
||||
result string
|
||||
out string
|
||||
}{
|
||||
{"no echo", "func main() {", "\n\tfmt.Println()", "\n\tfmt.Println()"},
|
||||
{"full echo", "func main() {", "func main() {\n\tfmt.Println()", "\n\tfmt.Println()"},
|
||||
{"partial tail echo", "hello world", "world new stuff", " new stuff"},
|
||||
{"empty prefix", "", "some result", "some result"},
|
||||
{"empty result", "prefix", "", ""},
|
||||
{"long prefix short echo", "a]" + string(make([]byte, 300)), "result", "result"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := stripPrefixEcho(tt.prefix, tt.result)
|
||||
if got != tt.out {
|
||||
t.Errorf("stripPrefixEcho(%q, %q) = %q, want %q", tt.prefix, tt.result, got, tt.out)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
209
main.go
209
main.go
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/crc32"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
|
|
@ -525,6 +526,206 @@ func (m *SessionManager) sendNotification(title, body string) {
|
|||
|
||||
// --- Shutdown ---
|
||||
|
||||
// --- Completion ---
|
||||
|
||||
// Complete performs a single-shot code completion using a dedicated copilot session.
|
||||
// It takes context about the cursor position and returns raw completion text.
|
||||
func (m *SessionManager) Complete(cwd, filePath, prefix, suffix, extraContext string) (string, *dbus.Error) {
|
||||
if cwd == "" {
|
||||
cwd, _ = os.Getwd()
|
||||
}
|
||||
cwd = paths.ExpandHome(os.ExpandEnv(cwd))
|
||||
|
||||
model := os.Getenv("OLLIE_COMPLETE_MODEL")
|
||||
backendName := os.Getenv("OLLIE_COMPLETE_BACKEND")
|
||||
if model == "" || backendName == "" {
|
||||
return "", dbus.MakeFailedError(fmt.Errorf("OLLIE_COMPLETE_MODEL and OLLIE_COMPLETE_BACKEND must be set"))
|
||||
}
|
||||
|
||||
// Truncate to context budget
|
||||
const prefixMax = 12000
|
||||
const suffixMax = 1000
|
||||
if len(prefix) > prefixMax {
|
||||
prefix = prefix[len(prefix)-prefixMax:]
|
||||
}
|
||||
if len(suffix) > suffixMax {
|
||||
suffix = suffix[:suffixMax]
|
||||
}
|
||||
|
||||
// Find or create a copilot session keyed by cwd
|
||||
sessID := m.findOrCreateCopilot(cwd, backendName, model)
|
||||
if sessID == "" {
|
||||
return "", dbus.MakeFailedError(fmt.Errorf("failed to create copilot session"))
|
||||
}
|
||||
|
||||
m.mu.RLock()
|
||||
sess, ok := m.sessions[sessID]
|
||||
m.mu.RUnlock()
|
||||
if !ok {
|
||||
return "", dbus.MakeFailedError(fmt.Errorf("copilot session disappeared"))
|
||||
}
|
||||
|
||||
// Build FIM prompt
|
||||
fileHint := ""
|
||||
if filePath != "" {
|
||||
fileHint = " in " + filePath
|
||||
}
|
||||
contextBlock := ""
|
||||
if extraContext != "" {
|
||||
contextBlock = "\n" + extraContext
|
||||
}
|
||||
prompt := fmt.Sprintf(`Implement the code at the cursor%s. The prefix ends at the point where new code is needed. Write the implementation — do not echo stubs, TODOs, or placeholder returns from the prefix. Output ONLY raw code. No reasoning, no shell commands, no explanations, no markdown fences, no backticks, no preamble. Your entire response must be valid code that can be inserted directly into the file.
|
||||
%s
|
||||
<prefix>
|
||||
%s
|
||||
</prefix>
|
||||
<suffix>
|
||||
%s
|
||||
</suffix>`, fileHint, contextBlock, prefix, suffix)
|
||||
|
||||
// Submit and wait for completion
|
||||
sess.core.Submit(sess.ctx, "/clear")
|
||||
sess.core.Submit(sess.ctx, prompt)
|
||||
|
||||
// Wait for idle
|
||||
for {
|
||||
state := sess.core.State()
|
||||
if state == "idle" {
|
||||
break
|
||||
}
|
||||
next, ok := sess.core.WaitChange(sess.ctx, agent.WatchState, state)
|
||||
if !ok {
|
||||
return "", dbus.MakeFailedError(fmt.Errorf("copilot session cancelled"))
|
||||
}
|
||||
if next == "idle" {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
result := sess.core.Reply()
|
||||
|
||||
// Strip markdown fences and leaked XML tags
|
||||
result = stripCompletionNoise(result)
|
||||
|
||||
// Strip prefix echo
|
||||
result = stripPrefixEcho(prefix, result)
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// findOrCreateCopilot finds an existing copilot session for cwd or creates one.
|
||||
func (m *SessionManager) findOrCreateCopilot(cwd, backendName, modelName string) string {
|
||||
// Session ID is deterministic based on cwd
|
||||
sessID := fmt.Sprintf("%d-copilot", crc32Str(cwd))
|
||||
|
||||
m.mu.RLock()
|
||||
_, exists := m.sessions[sessID]
|
||||
m.mu.RUnlock()
|
||||
|
||||
if exists {
|
||||
return sessID
|
||||
}
|
||||
|
||||
// Create a new copilot session
|
||||
agentsDir := paths.CfgDir() + "/agents"
|
||||
var cfg *config.Config
|
||||
cfgPath := agent.AgentConfigPath(agentsDir, "copilot")
|
||||
if f, err := os.Open(cfgPath); err == nil {
|
||||
cfg, _ = config.Load(f)
|
||||
f.Close()
|
||||
}
|
||||
|
||||
be, err := backend.NewWithName(backendName)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
be.SetModel(modelName)
|
||||
|
||||
sessionsDir := paths.DataDir() + "/sessions"
|
||||
os.MkdirAll(sessionsDir, 0700)
|
||||
|
||||
newDisp := tools.NewDispatcherFunc(map[string]func() tools.Server{
|
||||
"execute": execute.Decl(cwd),
|
||||
})
|
||||
rt := agent.BuildRuntime(cfg, newDisp(), cwd, []string{"OLLIE_SESSION_ID=" + sessID})
|
||||
|
||||
core := agent.NewAgentCore(agent.AgentCoreConfig{
|
||||
Backend: be,
|
||||
AgentName: "copilot",
|
||||
AgentsDir: agentsDir,
|
||||
SessionsDir: sessionsDir,
|
||||
SessionID: sessID,
|
||||
CWD: cwd,
|
||||
Runtime: rt,
|
||||
NewDispatcher: newDisp,
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
sess := &managedSession{
|
||||
core: core,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
id: sessID,
|
||||
agent: "copilot",
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
// Double-check after lock
|
||||
if _, exists := m.sessions[sessID]; exists {
|
||||
m.mu.Unlock()
|
||||
core.Close()
|
||||
cancel()
|
||||
return sessID
|
||||
}
|
||||
m.sessions[sessID] = sess
|
||||
m.mu.Unlock()
|
||||
|
||||
return sessID
|
||||
}
|
||||
|
||||
// stripCompletionNoise removes markdown fences and leaked XML tags.
|
||||
func stripCompletionNoise(s string) string {
|
||||
var lines []string
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
// Skip markdown code fences
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
continue
|
||||
}
|
||||
// Skip leaked XML tags
|
||||
if trimmed == "<prefix>" || trimmed == "</prefix>" ||
|
||||
trimmed == "<suffix>" || trimmed == "</suffix>" {
|
||||
continue
|
||||
}
|
||||
// Skip info lines
|
||||
if strings.HasPrefix(trimmed, ":: ") {
|
||||
continue
|
||||
}
|
||||
lines = append(lines, line)
|
||||
}
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
// stripPrefixEcho removes echoed prefix from the beginning of the result.
|
||||
func stripPrefixEcho(prefix, result string) string {
|
||||
tailMax := 200
|
||||
if len(prefix) < tailMax {
|
||||
tailMax = len(prefix)
|
||||
}
|
||||
for i := tailMax; i > 0; i-- {
|
||||
tail := prefix[len(prefix)-i:]
|
||||
if strings.HasPrefix(result, tail) {
|
||||
return result[len(tail):]
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// crc32Str returns a CRC32 checksum of a string as uint32.
|
||||
func crc32Str(s string) uint32 {
|
||||
return crc32.ChecksumIEEE([]byte(s))
|
||||
}
|
||||
|
||||
func (m *SessionManager) Shutdown() {
|
||||
m.mu.Lock()
|
||||
for id, sess := range m.sessions {
|
||||
|
|
@ -630,6 +831,14 @@ const introspectXML = `<node>
|
|||
<arg name="prompt" type="s" direction="in"/>
|
||||
<arg name="success" type="b" direction="out"/>
|
||||
</method>
|
||||
<method name="Complete">
|
||||
<arg name="cwd" type="s" direction="in"/>
|
||||
<arg name="file" type="s" direction="in"/>
|
||||
<arg name="prefix" type="s" direction="in"/>
|
||||
<arg name="suffix" type="s" direction="in"/>
|
||||
<arg name="context" type="s" direction="in"/>
|
||||
<arg name="result" type="s" direction="out"/>
|
||||
</method>
|
||||
<signal name="SessionCreated">
|
||||
<arg name="session_id" type="s"/>
|
||||
</signal>
|
||||
|
|
|
|||
Reference in New Issue