Design an AI Coding Assistant

Design an AI pair programmer that lives inside the IDE. As a developer types, grey "ghost text" suggests the rest of the line or the next few lines, and Tab accepts it. A side panel lets the developer chat about the codebase: explain this function, write a test, refactor this file.

Inline suggestions are only useful if they appear before the developer types the next character, so the latency budget is a few hundred milliseconds end to end, on every keystroke pause, for a million developers. Suggestions are much better when the model sees relevant code from other files, but the model can only take a limited prompt. Many customers are companies whose source code must not be stored or used for training.

In 45 minutes, cover the completion path and its latency budget, how relevant repository context is selected, caching, the chat path, privacy guarantees, and the telemetry you need to know whether suggestions help.

Sketch on the HLD board, write your answer in the Design doc, and size it with the Estimator. Start the mock interview from the timer for a 45-minute run.

Design Path · AI & ML Systems · problem 55 of 70

Ask before you design. Try answering each yourself, then compare.

1. How many users and how often do they trigger completions?
   Typical answer

   ~1 M daily developers; after client-side debouncing, ~500 completion requests per developer per day, concentrated in working hours.

2. What latency target for inline completions?
   Typical answer

   p50 under ~300 ms and p95 under ~600 ms from keystroke pause to ghost text. Chat can take seconds but must stream.

3. Which models?
   Typical answer

   A small self-hosted code model (1–7B parameters, trained for fill-in-the-middle) for completions; a larger hosted or self-hosted model for chat.

4. What code may leave the developer's machine?
   Typical answer

   Snippets needed for a request may be sent over TLS; business customers require zero retention, no training on their code, and region pinning.

5. Do we index whole repositories server-side?
   Typical answer

   Yes for business tenants that opt in (for chat and cross-file context); otherwise context comes only from the local workspace.

6. How do we measure success?
   Typical answer

   Acceptance rate, and how much accepted code is still present a few minutes later, plus latency.

Requirements

Functional

• Inline completions using code before and after the cursor (fill-in-the-middle).
• Chat about the code with references to files and symbols.
• Cross-file context from the workspace and, optionally, a server-side repo index.
• Admin controls: content exclusions, retention policy, region.
• Telemetry on suggestions shown, accepted and retained.

Non-functional

• Completion p50 < 300 ms, p95 < 600 ms; chat first token < 1.5 s.
• ~500 M completion requests/day, ~23k/s peak.
• 99.9% availability; failures degrade silently (no ghost text) rather than blocking typing.
• Zero code retention for business tenants; tenant isolation of indexes.

Out of scope

• Training the models.
• Autonomous agents that run commands or open pull requests.
• IDE plugin UI details.

Back-of-the-envelope estimates

• Requests: 1 M devs × 500 = 500 M/day ≈ 5,800/s average; working hours concentrate it to ~23k/s at peak.
• Prompt size (assumption): ~2k tokens (prefix, suffix and cross-file snippets); output ~30 tokens.
• Prefill at peak: 23k × 2k ≈ 46 M tokens/s. Consecutive requests share most of the prompt, so with session-sticky KV prefix reuse only ~25% (~12 M tokens/s) needs computing.
• GPUs (assumptions: small code model at ~150k prefill and ~10k decode tokens/s per GPU): ~80 GPUs for prefill and ~70 for decode (23k × 30 = 0.7 M tokens/s), about 150 busy GPUs; spread over 3 regions with headroom, ~400 GPUs.
• Cancellations: roughly half of requests are superseded by a later keystroke; cancelling them early is pure savings.
• Telemetry: ~1 KB event per request ≈ 500 GB/day.

The latency budget, not throughput, drives the design: small model, regional serving close to users, prefix reuse, and aggressive cancellation.

API

POST /v1/completions   {requestId, file:{path, lang}, prefix, suffix,
                        context:[{path, text, source}], maxTokens: 64}
     -> stream {text, finish}            (HTTP/2; client may RST to cancel)
POST /v1/chat          {conversationId, messages, references:[{path, range}]}
     -> SSE deltas with file/line citations
POST /v1/telemetry     {requestId, event: shown|accepted|rejected|retained_30s, chars}
PUT  /v1/admin/policy  {exclusions[], retention: none, region}

The client assembles most of the context itself and sends it with the request, because it already knows the open tabs, recent edits and language-server symbols, and that avoids a server round trip on the hot path.

Data model

Completion requests are processed in memory and not stored for zero-retention tenants. The server-side code index holds, per tenant and repository branch, chunk(tenant_id, repo, path, symbol, start_line, end_line, content_hash, vector) in a vector store with a parallel symbol table (symbol → definition location) for exact lookups; it is partitioned by tenant so one customer's queries never scan another's data, and encrypted with per-tenant keys.

Telemetry events (request_id, user_hash, lang, latency_ms, shown, accepted, retained_chars, model_version, experiment) contain no code by default and land in a columnar warehouse for dashboards and A/B analysis. Admin policy (exclusions, retention, region) is a small relational table cached at the edge.

High-level design

Completion path: the IDE extension waits for a short typing pause (~50–75 ms), first checks its local cache (the developer may be typing exactly what was suggested), then builds the request: code before and after the cursor plus a few ranked snippets from open tabs, recent edits and definitions resolved by the language server. It sends this over a persistent HTTP/2 connection to the nearest regional edge gateway, which authenticates, applies tenant policy and checks the server completion cache. The completion service adds server-side context from the context engine when the tenant has a repo index and the time budget allows, formats a fill-in-the-middle prompt, and routes to a completion model replica, sticky per developer session so the previous request's KV cache prefix is reused. Tokens stream back and the client stops early at a natural boundary such as end of block. A newer keystroke cancels the in-flight request, which frees the GPU slot.

Chat path: the chat service asks the context engine to retrieve relevant chunks and symbols from the code index, which the repo indexer keeps fresh from code host push events, and streams an answer from the larger chat LLM with file and line references. Both paths send outcome events to telemetry, aggregated in the warehouse.

Open the HLD board and press Reference to see the architecture diagram.

Deep dives

Spending a 300 ms budget

Rough budget: debounce 50–75 ms, network round trip 20–60 ms to a regional edge, queueing under 10 ms, prefill of ~2k tokens on a small model 20–40 ms (less with a reused prefix), and decode of ~20–30 tokens at a few ms each. That only works with a small model, regional GPUs, warm persistent connections, and short outputs: stop at the end of the line or block instead of generating 128 tokens. Cancellation is essential since about half of requests are stale before they finish. Batching on the server must use small maximum batch delays, because waiting 20 ms to fill a batch spends a big share of the budget. Speculative decoding with a tiny draft model helps when outputs are longer.

Choosing context that fits

The highest-value context is local and cheap: text before and after the cursor, then definitions of symbols near the cursor (from the language server), recently edited regions, and similar snippets from open tabs found with a quick token-overlap score. Each source gets a token budget and candidates are ranked and packed greedily so the most useful lines win. Server-side retrieval from an embedding index adds value for chat and for repos where relevant code is not open, but it costs a round trip, so for completions it is used only when it fits in the budget or is precomputed per file. Order matters: put the most relevant snippets closest to the cursor text, and keep the file prefix stable across requests so prefix reuse works.

Caching at three levels

Client: an LRU of recent suggestions; if the developer types characters that match the suggestion, show the remainder with no request. This also covers backspace-and-retype. Server: an exact result cache keyed by a hash of the normalized prompt has a modest hit rate (~10–20%) because prompts are so specific, but it is nearly free. GPU: the biggest win is reusing attention KV state for a shared prefix. Successive requests from one session differ only by the last few characters, so session-affine routing (consistent hash of session id, with bounded load) lets the server compute only the new tail. That can cut prefill work by well over half.

Privacy and telemetry

Business tenants get zero retention: prompts and outputs live only in memory for the request, logs carry ids and sizes but no code, and providers behind the chat path are contracted for no storage or training. Admins can exclude paths or repositories; the client enforces exclusions before sending anything, and the server enforces them again. Secrets are detected and stripped from prompts and suggestions. Indexes are per tenant, encrypted with tenant keys and pinned to a region. Telemetry measures acceptance rate, characters retained after 30 seconds and 5 minutes, latency and cancellations, broken down by language and model version; experiments compare models on these metrics before rollout.

Failure modes

• If a completion replica or region is slow, the client times out at ~800 ms and shows nothing; the gateway routes subsequent requests to another replica or region.
• If the code index is stale or down, completions still work from client-sent context and chat falls back to open files only.
• If GPU capacity runs short, the service shortens max output and context, and drops speculative requests before rejecting.
• If an exclusion policy fails to load, the client fails closed for that workspace and sends no code.

Core algorithm in code

import re
from collections import OrderedDict


def idents(text):
    return set(re.findall(r"[A-Za-z_][A-Za-z0-9_]{2,}", text))


def tokens(text):
    return max(1, len(text) // 4)          # rough rule of thumb: ~4 chars per token


SOURCE_BOOST = {"definition": 1.5, "recent_edit": 1.3, "open_tab": 1.0, "repo_search": 0.8}


def pack_context(cursor_prefix, snippets, budget_tokens):
    near = idents(cursor_prefix[-600:])
    scored = []
    for s in snippets:
        ids = idents(s["text"])
        overlap = len(near & ids) / (len(ids) or 1)
        scored.append((overlap * SOURCE_BOOST[s["source"]], s))
    scored.sort(key=lambda x: -x[0])
    picked, used = [], 0
    for score, s in scored:
        cost = tokens(s["text"])
        if score > 0 and used + cost <= budget_tokens:
            picked.append((round(score, 2), s["path"]))
            used += cost
    return picked, used


class CompletionCache:
    def __init__(self, capacity=256):
        self.cap, self.items = capacity, OrderedDict()

    def put(self, prefix, completion):
        self.items[prefix] = completion
        self.items.move_to_end(prefix)
        if len(self.items) > self.cap:
            self.items.popitem(last=False)

    def get(self, prefix):
        # exact hit, or the user typed part of an earlier suggestion
        for p in reversed(self.items):
            full = p + self.items[p]
            if prefix.startswith(p) and full.startswith(prefix) and len(full) > len(prefix):
                self.items.move_to_end(p)
                return full[len(prefix):]
        return None


cursor = """from billing.invoice import Invoice, apply_discount

def close_month(customer):
    inv = Invoice.for_customer(customer)
    total = apply_discount(inv.subtotal, customer."""
snippets = [
    {"path": "billing/invoice.py:Invoice", "source": "definition",
     "text": "class Invoice:\n    subtotal: Decimal\n    def for_customer(cls, customer): ..."},
    {"path": "billing/invoice.py:apply_discount", "source": "definition",
     "text": "def apply_discount(amount, tier):\n    return amount * DISCOUNTS[tier]"},
    {"path": "customers/model.py", "source": "open_tab",
     "text": "class Customer:\n    tier: str\n    region: str\n    customer_id: int"},
    {"path": "ui/theme.css", "source": "open_tab", "text": "body { color: #222; margin: 0 }"},
    {"path": "billing/tests.py", "source": "recent_edit",
     "text": "def test_close_month():\n    inv = Invoice(subtotal=100)\n    assert apply_discount(100, 'gold') == 90"},
]
picked, used = pack_context(cursor, snippets, budget_tokens=60)
print("context picked (score, path):", picked, "| tokens used:", used)

cache = CompletionCache()
cache.put("total = apply_discount(inv.subtotal, customer.", "tier)")
for typed in ("total = apply_discount(inv.subtotal, customer.",
              "total = apply_discount(inv.subtotal, customer.ti",
              "total = apply_discount(inv.subtotal, customer.re"):
    print("%-50r -> %r" % (typed[-20:], cache.get(typed)))

Trade-offs

• A small completion model meets the latency budget but costs suggestion quality on harder multi-line edits.
• More cross-file context improves suggestions but costs prefill time and pushes latency past the budget.
• Session-sticky routing enables KV prefix reuse but risks hot replicas and complicates failover.
• Server-side repo indexing improves chat answers but means storing customer code, which some tenants forbid.
• Rich telemetry speeds up model improvement but conflicts with strict privacy promises unless it carries no code.

What a strong answer shows

• Breaks down the latency budget and designs to it (small model, regional serving, cancellation).
• Explains how context is selected and packed under a token budget.
• Uses prefix reuse and client-side caching, not just a generic result cache.
• Separates the completion and chat paths with different models and SLOs.
• Makes concrete privacy guarantees (retention, exclusions, tenant isolation).
• Defines meaningful metrics beyond raw acceptance rate.

Mistakes to avoid

• Sending every keystroke to a large model with no debounce or cancellation.
• Retrieving from a repo-wide vector index on the completion hot path regardless of latency.
• Treating it as one service with one model for both completion and chat.
• Logging prompts for debugging, which breaks zero-retention promises.
• Measuring only acceptance rate, which rewards short, trivial suggestions.

Follow-ups to expect

• How would you support offline or air-gapped customers? (Smaller quantized model on device or on-prem GPUs.)
• How would you suggest multi-file edits? (Chat path plans edits; IDE shows a diff for review.)
• How do you avoid suggesting licensed code verbatim? (Filter outputs against an index of known public code.)
• How would you roll out a new completion model? (Shadow, then A/B on retained characters and latency.)

Practice next